chore: import upstream snapshot with attribution

python/ray/air/BUILD.bazel

@@ -0,0 +1,305 @@
+load("@rules_python//python:defs.bzl", "py_library", "py_test")
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    name = "py_doctest[air]",
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    files = glob(
+        ["**/*.py"],
+        exclude = glob([
+            "examples/**/*",
+            "tests/**/*",
+            "callbacks/*.py",
+        ]) + ["integrations/wandb.py"],
+    ),
+    tags = ["team:ml"],
+)
+
+py_library(
+    name = "conftest",
+    srcs = ["tests/conftest.py"],
+)
+
+# --------------------------------------------------------------------
+# Tests from the python/ray/air/tests directory.
+# Covers all tests starting with `test_`.
+# Please keep these sorted alphabetically.
+# --------------------------------------------------------------------
+
+py_test(
+    name = "test_air_usage",
+    size = "small",
+    srcs = ["tests/test_air_usage.py"],
+    # NOTE: This tests Train V1 telemetry.
+    env = {"RAY_TRAIN_V2_ENABLED": "0"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_new_dataset_config",
+    size = "large",
+    srcs = ["tests/test_new_dataset_config.py"],
+    # NOTE: Relevant tests moved to train/v2/tests/test_data_integration.py
+    env = {"RAY_TRAIN_V2_ENABLED": "0"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_experiment_restore",
+    size = "large",
+    srcs = [
+        "tests/_test_experiment_restore_run.py",
+        "tests/test_experiment_restore.py",
+    ],
+    # NOTE: This tests Tune and Train V1 restoration.
+    env = {"RAY_TRAIN_V2_ENABLED": "0"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_errors",
+    size = "medium",
+    srcs = ["tests/test_errors.py"],
+    # NOTE: This tests Tune (Train V1) error propagation logic.
+    env = {"RAY_TRAIN_V2_ENABLED": "0"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_integration_comet",
+    size = "small",
+    srcs = ["tests/test_integration_comet.py"],
+    # NOTE: This tests the Tune Comet callback.
+    env = {"RAY_TRAIN_V2_ENABLED": "0"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_integration_wandb",
+    size = "small",
+    srcs = ["tests/test_integration_wandb.py"],
+    # NOTE: This tests the Tune wandb callback.
+    env = {"RAY_TRAIN_V2_ENABLED": "0"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_integration_mlflow",
+    size = "medium",
+    srcs = ["tests/test_integration_mlflow.py"],
+    # NOTE: This tests the Tune mlflow callback.
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_keras_callback",
+    size = "medium",
+    srcs = ["tests/test_keras_callback.py"],
+    env = {
+        "RAY_TRAIN_V2_ENABLED": "1",
+        "TF_USE_LEGACY_KERAS": "1",
+    },
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_remote_storage_hdfs",
+    size = "small",
+    srcs = ["tests/test_remote_storage_hdfs.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "hdfs",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_tracebacks",
+    size = "small",
+    srcs = ["tests/test_tracebacks.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_utils",
+    size = "small",
+    srcs = ["tests/test_utils.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+# --------------------------------------------------------------------
+# Tests from the python/ray/air/tests/execution directory.
+# Covers all tests starting with `test_`.
+# Please keep these sorted alphabetically.
+# TODO: Move this to Tune.
+# --------------------------------------------------------------------
+
+py_test(
+    name = "test_barrier",
+    size = "small",
+    srcs = ["tests/execution/test_barrier.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_e2e_train_flow",
+    size = "medium",
+    srcs = ["tests/execution/test_e2e_train_flow.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_e2e_tune_flow",
+    size = "medium",
+    srcs = ["tests/execution/test_e2e_tune_flow.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_event_manager",
+    size = "medium",
+    srcs = ["tests/execution/test_event_manager.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_resource_manager_fixed",
+    size = "small",
+    srcs = ["tests/execution/test_resource_manager_fixed.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_resource_manager_placement_group",
+    size = "medium",
+    srcs = ["tests/execution/test_resource_manager_placement_group.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_resource_request",
+    size = "small",
+    srcs = ["tests/execution/test_resource_request.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_tracked_actor",
+    size = "small",
+    srcs = ["tests/execution/test_tracked_actor.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+py_test(
+    name = "test_tracked_actor_task",
+    size = "small",
+    srcs = ["tests/execution/test_tracked_actor_task.py"],
+    env = {"RAY_TRAIN_V2_ENABLED": "1"},
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = [":ml_lib"],
+)
+
+# This is a dummy test dependency that causes the above tests to be
+# re-run if any of these files changes.
+py_library(
+    name = "ml_lib",
+    srcs = glob(
+        ["**/*.py"],
+        exclude = ["tests/*.py"],
+    ),
+    visibility = [
+        "//python/ray/air:__pkg__",
+        "//python/ray/air:__subpackages__",
+        "//python/ray/train:__pkg__",
+        "//python/ray/train:__subpackages__",
+        "//release:__pkg__",
+    ],
+)

python/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",
+]

python/ray/air/_internal/config.py

@@ -0,0 +1,52 @@
+import dataclasses
+from typing import Iterable
+
+
+def _validate_allowed_keys_exist(
+    dataclass_name: str, data_dict: dict, allowed_keys: set
+):
+    keys_not_in_dict = allowed_keys.difference(data_dict)
+    if keys_not_in_dict:
+        raise ValueError(
+            f"Key(s) {sorted(keys_not_in_dict)} are not present in {dataclass_name}. "
+            "Remove them from `allowed_keys`. "
+            f"Valid keys: {sorted(data_dict.keys())}"
+        )
+
+
+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__()
+    default_data_dict = default_data.__dict__
+
+    allowed_keys = set(allowed_keys)
+    _validate_allowed_keys_exist(
+        dataclass.__class__.__name__, default_data_dict, allowed_keys
+    )
+
+    # These keys should not have been updated in the `dataclass` object
+    prohibited_keys = set(default_data_dict) - allowed_keys
+    dataclass_dict = dataclass.__dict__
+
+    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."
+        )

python/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,
+]

python/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()

python/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()

python/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()

python/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()

python/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"""
+        ...

python/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_default_system_temp_dir()).resolve()
+        self._lock_dir = temp_dir / RAY_LOCKFILE_DIR
+        self._path_hash = hashlib.sha256(
+            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)

python/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)

python/ray/air/_internal/mlflow.py

@@ -0,0 +1,346 @@
+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,
+    ) -> None:
+        """
+        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.
+            tracking_token: Tracking token used to authenticate with MLflow.
+            artifact_location: The location to store run artifacts.
+                If not provided, MLFlow picks an appropriate default.
+                Ignored if experiment already exists.
+            create_experiment_if_not_exists: Whether to create an
+                experiment with the provided name if it does not already
+                exist. Defaults to True.
+
+        Raises:
+             ValueError: ``experiment_id`` and ``experiment_name`` are both ``None``.
+        """
+        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:
+            run_name: Name of the new MLflow run to create.
+            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: 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: int, 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:
+            step: Step at which the metrics are logged.
+            metrics_to_log: Dictionary of metrics to log.
+            run_id: 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: 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: Optional[str] = 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: The status to set when terminating the run.
+            run_id: 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)

python/ray/air/_internal/tensorflow_utils.py

@@ -0,0 +1,73 @@
+from typing import Dict, Optional, Union
+
+import numpy as np
+import tensorflow as tf
+
+from ray.air.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed
+
+
+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:
+        ndarrays: A (dict of) NumPy ndarray(s) that we wish to convert to a TensorFlow
+            Tensor.
+        dtypes: 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

python/ray/air/_internal/torch_utils.py

@@ -0,0 +1,105 @@
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+
+from ray.air._internal.device_manager import get_torch_device_manager_by_context
+
+
+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 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.
+
+    Returns:
+        A new state-dict with the prefix stripped from the keys.
+    """
+    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

python/ray/air/_internal/uri_utils.py

@@ -0,0 +1,106 @@
+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'
+
+        Args:
+            subpath: The subpath to strip from the end of this URI.
+
+        Returns:
+            A new URI with the subpath stripped from the end.
+        """
+        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)

python/ray/air/_internal/usage.py

@@ -0,0 +1,277 @@
+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: object, 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.
+
+    Args:
+        callbacks: List of callbacks supplied by the user. May be ``None``.
+
+    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)

python/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

python/ray/air/config.py

@@ -0,0 +1,723 @@
+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: Any, *, 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.
+
+            Accepts the same resource keys that Ray uses for scheduling tasks
+            and actors (see :ref:`Resources <core-resources>`):
+
+            - ``"CPU"``: number of logical CPUs.
+            - ``"GPU"``: number of logical GPUs.
+            - ``"memory"``: heap memory reserved on the node, in bytes
+              (for example, ``"memory": 1e9`` reserves 1 GB).
+            - Any :ref:`custom resource <custom-resources>` name configured on
+              your cluster (for example, ``"TPU": 1``, ``"special_hardware": 1``).
+
+            Keys are case-sensitive: use ``"CPU"`` and ``"GPU"`` (uppercase),
+            and ``"memory"`` (lowercase).
+        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.
+
+            Accepts the same resource keys as ``trainer_resources``:
+
+            - ``"CPU"``: number of logical CPUs per worker.
+            - ``"GPU"``: number of logical GPUs per worker. Prefer setting
+              ``use_gpu=True`` (which reserves 1 GPU per worker) and only
+              override this key when you need a different per-worker count.
+            - ``"memory"``: heap memory reserved per worker, in bytes
+              (for example, ``"memory": 1e9`` reserves 1 GB per worker).
+            - Any :ref:`custom resource <custom-resources>` name configured on
+              your cluster (for example, ``"TPU": 1``, ``"special_hardware": 1``).
+
+            Keys are case-sensitive: use ``"CPU"`` and ``"GPU"`` (uppercase),
+            and ``"memory"`` (lowercase).
+        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,
+            ),
+        )

python/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",
+}

python/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"]
+]

python/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",
+]

python/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"]

python/ray/air/execution/_internal/actor_manager.py

@@ -0,0 +1,906 @@
+import logging
+import random
+import time
+import uuid
+from collections import Counter, defaultdict
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
+
+import ray
+from ray.air.execution._internal.event_manager import RayEventManager
+from ray.air.execution._internal.tracked_actor import TrackedActor
+from ray.air.execution._internal.tracked_actor_task import TrackedActorTask
+from ray.air.execution.resources import (
+    AcquiredResources,
+    ResourceManager,
+    ResourceRequest,
+)
+from ray.exceptions import RayActorError, RayTaskError
+
+logger = logging.getLogger(__name__)
+
+
+class RayActorManager:
+    """Management class for Ray actors and actor tasks.
+
+    This class provides an event-based management interface for actors, and
+    actor tasks.
+
+    The manager can be used to start actors, stop actors, and schedule and
+    track task futures on these actors.
+    The manager will then invoke callbacks related to the tracked entities.
+
+    For instance, when an actor is added with
+    :meth:`add_actor() <RayActorManager.add_actor>`,
+    a :ref:`TrackedActor <ray.air.execution._internal.tracked_actor.TrackedActor`
+    object is returned. An ``on_start`` callback can be specified that is invoked
+    once the actor successfully started. Similarly, ``on_stop`` and ``on_error``
+    can be used to specify callbacks relating to the graceful or ungraceful
+    end of an actor's lifetime.
+
+    When scheduling an actor task using
+    :meth:`schedule_actor_task()
+    <ray.air.execution._internal.actor_manager.RayActorManager.schedule_actor_task>`,
+    an ``on_result`` callback can be specified that is invoked when the task
+    successfully resolves, and an ``on_error`` callback will resolve when the
+    task fails.
+
+    The RayActorManager does not implement any true asynchronous processing. Control
+    has to be explicitly yielded to the event manager via :meth:`RayActorManager.next`.
+    Callbacks will only be invoked when control is with the RayActorManager, and
+    callbacks will always be executed sequentially in order of arriving events.
+
+    Args:
+        resource_manager: Resource manager used to request resources for the actors.
+
+    Example:
+
+        .. code-block:: python
+
+            from ray.air.execution import ResourceRequest
+            from ray.air.execution._internal import RayActorManager
+
+            actor_manager = RayActorManager()
+
+            # Request an actor
+            tracked_actor = actor_manager.add_actor(
+                ActorClass,
+                kwargs={},
+                resource_request=ResourceRequest([{"CPU": 1}]),
+                on_start=actor_start_callback,
+                on_stop=actor_stop_callback,
+                on_error=actor_error_callback
+            )
+
+            # Yield control to event manager to start actor
+            actor_manager.next()
+
+            # Start task on the actor (ActorClass.foo.remote())
+            tracked_actor_task = actor_manager.schedule_actor_task(
+                tracked_actor,
+                method_name="foo",
+                on_result=task_result_callback,
+                on_error=task_error_callback
+            )
+
+            # Again yield control to event manager to process task futures
+            actor_manager.wait()
+
+    """
+
+    def __init__(self, resource_manager: ResourceManager):
+        self._resource_manager: ResourceManager = resource_manager
+
+        self._actor_state_events = RayEventManager()
+        self._actor_task_events = RayEventManager()
+
+        # ---
+        # Tracked actor futures.
+
+        # This maps TrackedActor objects to their futures. We use this to see if an
+        # actor has any futures scheduled and to remove them when we terminate an actor.
+
+        # Actors to actor task futures
+        self._tracked_actors_to_task_futures: Dict[
+            TrackedActor, Set[ray.ObjectRef]
+        ] = defaultdict(set)
+
+        # Actors to actor state futures (start/terminate)
+        self._tracked_actors_to_state_futures: Dict[
+            TrackedActor, Set[ray.ObjectRef]
+        ] = defaultdict(set)
+
+        # ---
+        # Pending actors.
+        # We use three dicts for actors that are requested but not yet started.
+
+        # This dict keeps a list of actors associated with each resource request.
+        # We use this to start actors in the correct order when their resources
+        # become available.
+        self._resource_request_to_pending_actors: Dict[
+            ResourceRequest, List[TrackedActor]
+        ] = defaultdict(list)
+
+        # This dict stores the actor class, kwargs, and resource request of
+        # pending actors. Once the resources are available, we start the remote
+        # actor class with its args. We need the resource request to cancel it
+        # if needed.
+        self._pending_actors_to_attrs: Dict[
+            TrackedActor, Tuple[Type, Dict[str, Any], ResourceRequest]
+        ] = {}
+
+        # This dict keeps track of cached actor tasks. We can't schedule actor
+        # tasks before the actor is actually scheduled/live. So when the caller
+        # tries to schedule a task, we cache it here, and schedule it once the
+        # actor is started.
+        self._pending_actors_to_enqueued_actor_tasks: Dict[
+            TrackedActor, List[Tuple[TrackedActorTask, str, Tuple[Any], Dict[str, Any]]]
+        ] = defaultdict(list)
+
+        # ---
+        # Live actors.
+        # We keep one dict for actors that are currently running and a set of
+        # actors that we should forcefully kill.
+
+        # This dict associates the TrackedActor object with the Ray actor handle
+        # and the resources associated to the actor. We use it to schedule the
+        # actual ray tasks, and to return the resources when the actor stopped.
+        self._live_actors_to_ray_actors_resources: Dict[
+            TrackedActor, Tuple[ray.actor.ActorHandle, AcquiredResources]
+        ] = {}
+        self._live_resource_cache: Optional[Dict[str, Any]] = None
+
+        # This dict contains all actors that should be killed (after calling
+        # `remove_actor()`). Kill requests will be handled in wait().
+        self._live_actors_to_kill: Set[TrackedActor] = set()
+
+        # Track failed actors
+        self._failed_actor_ids: Set[int] = set()
+
+    def next(self, timeout: Optional[Union[int, float]] = None) -> bool:
+        """Yield control to event manager to await the next event and invoke callbacks.
+
+        Calling this method will wait for up to ``timeout`` seconds for the next
+        event to arrive.
+
+        When events arrive, callbacks relating to the events will be
+        invoked. A timeout of ``None`` will block until the next event arrives.
+
+        Note:
+            If an actor task fails with a ``RayActorError``, this is one event,
+            but it may trigger _two_ `on_error` callbacks: One for the actor,
+            and one for the task.
+
+        Note:
+            The ``timeout`` argument is used for pure waiting time for events. It does
+            not include time spent on processing callbacks. Depending on the processing
+            time of the callbacks, it can take much longer for this function to
+            return than the specified timeout.
+
+        Args:
+            timeout: Timeout in seconds to wait for next event.
+
+        Returns:
+            True if at least one event was processed.
+
+        """
+        # First issue any pending forceful actor kills
+        actor_killed = self._try_kill_actor()
+
+        # We always try to start actors as this won't trigger an event callback
+        self._try_start_actors()
+
+        # If an actor was killed, this was our event, and we return.
+        if actor_killed:
+            return True
+
+        # Otherwise, collect all futures and await the next.
+        resource_futures = self._resource_manager.get_resource_futures()
+        actor_state_futures = self._actor_state_events.get_futures()
+        actor_task_futures = self._actor_task_events.get_futures()
+
+        # Shuffle state futures
+        shuffled_state_futures = list(actor_state_futures)
+        random.shuffle(shuffled_state_futures)
+
+        # Shuffle task futures
+        shuffled_task_futures = list(actor_task_futures)
+        random.shuffle(shuffled_task_futures)
+
+        # Prioritize resource futures over actor state over task futures
+        all_futures = resource_futures + shuffled_state_futures + shuffled_task_futures
+
+        start_wait = time.monotonic()
+        ready, _ = ray.wait(all_futures, num_returns=1, timeout=timeout)
+
+        if not ready:
+            return False
+
+        [future] = ready
+
+        if future in actor_state_futures:
+            self._actor_state_events.resolve_future(future)
+        elif future in actor_task_futures:
+            self._actor_task_events.resolve_future(future)
+        else:
+            self._handle_ready_resource_future()
+            # Ready resource futures don't count as one event as they don't trigger
+            # any callbacks. So we repeat until we hit anything that is not a resource
+            # future.
+            time_taken = time.monotonic() - start_wait
+            return self.next(
+                timeout=max(1e-9, timeout - time_taken) if timeout is not None else None
+            )
+
+        self._try_start_actors()
+        return True
+
+    def _actor_start_resolved(self, tracked_actor: TrackedActor, future: ray.ObjectRef):
+        """Callback to be invoked when actor started"""
+        self._tracked_actors_to_state_futures[tracked_actor].remove(future)
+
+        if tracked_actor._on_start:
+            tracked_actor._on_start(tracked_actor)
+
+    def _actor_stop_resolved(self, tracked_actor: TrackedActor):
+        """Callback to be invoked when actor stopped"""
+        self._cleanup_actor(tracked_actor=tracked_actor)
+
+        if tracked_actor._on_stop:
+            tracked_actor._on_stop(tracked_actor)
+
+    def _actor_start_failed(self, tracked_actor: TrackedActor, exception: Exception):
+        """Callback to be invoked when actor start/stop failed"""
+        self._failed_actor_ids.add(tracked_actor.actor_id)
+
+        self._cleanup_actor(tracked_actor=tracked_actor)
+
+        if tracked_actor._on_error:
+            tracked_actor._on_error(tracked_actor, exception)
+
+    def _actor_task_failed(
+        self, tracked_actor_task: TrackedActorTask, exception: Exception
+    ):
+        """Handle an actor task future that became ready.
+
+        - On actor error, trigger actor error callback AND error task error callback
+        - On task error, trigger actor task error callback
+        - On success, trigger actor task result callback
+        """
+        tracked_actor = tracked_actor_task._tracked_actor
+
+        if isinstance(exception, RayActorError):
+            self._failed_actor_ids.add(tracked_actor.actor_id)
+
+            # Clean up any references to the actor and its futures
+            self._cleanup_actor(tracked_actor=tracked_actor)
+
+            # Handle actor state callbacks
+            if tracked_actor._on_error:
+                tracked_actor._on_error(tracked_actor, exception)
+
+            # Then trigger actor task error callback
+            if tracked_actor_task._on_error:
+                tracked_actor_task._on_error(tracked_actor, exception)
+
+        elif isinstance(exception, RayTaskError):
+            # Otherwise only the task failed. Invoke callback
+            if tracked_actor_task._on_error:
+                tracked_actor_task._on_error(tracked_actor, exception)
+        else:
+            raise RuntimeError(
+                f"Caught unexpected exception: {exception}"
+            ) from exception
+
+    def _actor_task_resolved(self, tracked_actor_task: TrackedActorTask, result: Any):
+        tracked_actor = tracked_actor_task._tracked_actor
+
+        # Trigger actor task result callback
+        if tracked_actor_task._on_result:
+            tracked_actor_task._on_result(tracked_actor, result)
+
+    def _handle_ready_resource_future(self):
+        """Handle a resource future that became ready.
+
+        - Update state of the resource manager
+        - Try to start one actor
+        """
+        # Force resource manager to update internal state
+        self._resource_manager.update_state()
+        # We handle resource futures one by one, so only try to start 1 actor at a time
+        self._try_start_actors(max_actors=1)
+
+    def _try_start_actors(self, max_actors: Optional[int] = None) -> int:
+        """Try to start up to ``max_actors`` actors.
+
+        This function will iterate through all resource requests we collected for
+        pending actors. As long as a resource request can be fulfilled (resources
+        are available), we try to start as many actors as possible.
+
+        This will schedule a `Actor.__ray_ready__()` future which, once resolved,
+        will trigger the `TrackedActor.on_start` callback.
+        """
+        started_actors = 0
+
+        # Iterate through all resource requests
+        for resource_request in self._resource_request_to_pending_actors:
+            if max_actors is not None and started_actors >= max_actors:
+                break
+
+            # While we have resources ready and there are actors left to schedule
+            while (
+                self._resource_manager.has_resources_ready(resource_request)
+                and self._resource_request_to_pending_actors[resource_request]
+            ):
+                # Acquire resources for actor
+                acquired_resources = self._resource_manager.acquire_resources(
+                    resource_request
+                )
+                assert acquired_resources
+
+                # Get tracked actor to start
+                candidate_actors = self._resource_request_to_pending_actors[
+                    resource_request
+                ]
+                assert candidate_actors
+
+                tracked_actor = candidate_actors.pop(0)
+
+                # Get actor class and arguments
+                actor_cls, kwargs, _ = self._pending_actors_to_attrs.pop(tracked_actor)
+
+                if not isinstance(actor_cls, ray.actor.ActorClass):
+                    actor_cls = ray.remote(actor_cls)
+
+                # Associate to acquired resources
+                [remote_actor_cls] = acquired_resources.annotate_remote_entities(
+                    [actor_cls]
+                )
+
+                # Start Ray actor
+                actor = remote_actor_cls.remote(**kwargs)
+
+                # Track
+                self._live_actors_to_ray_actors_resources[tracked_actor] = (
+                    actor,
+                    acquired_resources,
+                )
+                self._live_resource_cache = None
+
+                # Schedule ready future
+                future = actor.__ray_ready__.remote()
+
+                self._tracked_actors_to_state_futures[tracked_actor].add(future)
+
+                # We need to create the callbacks in a function so tracked_actors
+                # are captured correctly.
+                def create_callbacks(
+                    tracked_actor: TrackedActor, future: ray.ObjectRef
+                ):
+                    def on_actor_start(result: Any):
+                        self._actor_start_resolved(
+                            tracked_actor=tracked_actor, future=future
+                        )
+
+                    def on_error(exception: Exception):
+                        self._actor_start_failed(
+                            tracked_actor=tracked_actor, exception=exception
+                        )
+
+                    return on_actor_start, on_error
+
+                on_actor_start, on_error = create_callbacks(
+                    tracked_actor=tracked_actor, future=future
+                )
+
+                self._actor_state_events.track_future(
+                    future=future,
+                    on_result=on_actor_start,
+                    on_error=on_error,
+                )
+
+                self._enqueue_cached_actor_tasks(tracked_actor=tracked_actor)
+
+                started_actors += 1
+
+        return started_actors
+
+    def _enqueue_cached_actor_tasks(self, tracked_actor: TrackedActor):
+        assert tracked_actor in self._live_actors_to_ray_actors_resources
+
+        # Enqueue cached futures
+        cached_tasks = self._pending_actors_to_enqueued_actor_tasks.pop(
+            tracked_actor, []
+        )
+        for tracked_actor_task, method_name, args, kwargs in cached_tasks:
+            self._schedule_tracked_actor_task(
+                tracked_actor_task=tracked_actor_task,
+                method_name=method_name,
+                args=args,
+                kwargs=kwargs,
+            )
+
+    def _try_kill_actor(self) -> bool:
+        """Try to kill actor scheduled for termination."""
+        if not self._live_actors_to_kill:
+            return False
+
+        tracked_actor = self._live_actors_to_kill.pop()
+
+        # Remove from tracked actors
+        (
+            ray_actor,
+            acquired_resources,
+        ) = self._live_actors_to_ray_actors_resources[tracked_actor]
+
+        # Hard kill if requested
+        ray.kill(ray_actor)
+
+        self._cleanup_actor_futures(tracked_actor)
+
+        self._actor_stop_resolved(tracked_actor)
+
+        return True
+
+    def _cleanup_actor(self, tracked_actor: TrackedActor):
+        self._cleanup_actor_futures(tracked_actor)
+
+        # Remove from tracked actors
+        (
+            ray_actor,
+            acquired_resources,
+        ) = self._live_actors_to_ray_actors_resources.pop(tracked_actor)
+        self._live_resource_cache = None
+
+        # Return resources
+        self._resource_manager.free_resources(acquired_resource=acquired_resources)
+
+    @property
+    def all_actors(self) -> List[TrackedActor]:
+        """Return all ``TrackedActor`` objects managed by this manager instance."""
+        return self.live_actors + self.pending_actors
+
+    @property
+    def live_actors(self) -> List[TrackedActor]:
+        """Return all ``TrackedActor`` objects that are currently alive."""
+        return list(self._live_actors_to_ray_actors_resources)
+
+    @property
+    def pending_actors(self) -> List[TrackedActor]:
+        """Return all ``TrackedActor`` objects that are currently pending."""
+        return list(self._pending_actors_to_attrs)
+
+    @property
+    def num_live_actors(self):
+        """Return number of started actors."""
+        return len(self.live_actors)
+
+    @property
+    def num_pending_actors(self) -> int:
+        """Return number of pending (not yet started) actors."""
+        return len(self.pending_actors)
+
+    @property
+    def num_total_actors(self):
+        """Return number of total actors."""
+        return len(self.all_actors)
+
+    @property
+    def num_actor_tasks(self):
+        """Return number of pending tasks"""
+        return self._actor_task_events.num_futures
+
+    def get_live_actors_resources(self):
+        if self._live_resource_cache:
+            return self._live_resource_cache
+
+        counter = Counter()
+        for _, acq in self._live_actors_to_ray_actors_resources.values():
+            for bdl in acq.resource_request.bundles:
+                counter.update(bdl)
+        self._live_resource_cache = dict(counter)
+        return self._live_resource_cache
+
+    def add_actor(
+        self,
+        cls: Union[Type, ray.actor.ActorClass],
+        kwargs: Dict[str, Any],
+        resource_request: ResourceRequest,
+        *,
+        on_start: Optional[Callable[[TrackedActor], None]] = None,
+        on_stop: Optional[Callable[[TrackedActor], None]] = None,
+        on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
+    ) -> TrackedActor:
+        """Add an actor to be tracked.
+
+        This method will request resources to start the actor. Once the resources
+        are available, the actor will be started and the
+        :meth:`TrackedActor.on_start
+        <ray.air.execution._internal.tracked_actor.TrackedActor.on_start>` callback
+        will be invoked.
+
+        Args:
+            cls: Actor class to schedule.
+            kwargs: Keyword arguments to pass to actor class on construction.
+            resource_request: Resources required to start the actor.
+            on_start: Callback to invoke when the actor started.
+            on_stop: Callback to invoke when the actor stopped.
+            on_error: Callback to invoke when the actor failed.
+
+        Returns:
+            Tracked actor object to reference actor in subsequent API calls.
+
+        """
+        tracked_actor = TrackedActor(
+            uuid.uuid4().int, on_start=on_start, on_stop=on_stop, on_error=on_error
+        )
+
+        self._pending_actors_to_attrs[tracked_actor] = cls, kwargs, resource_request
+        self._resource_request_to_pending_actors[resource_request].append(tracked_actor)
+
+        self._resource_manager.request_resources(resource_request=resource_request)
+
+        return tracked_actor
+
+    def remove_actor(
+        self,
+        tracked_actor: TrackedActor,
+        kill: bool = False,
+        stop_future: Optional[ray.ObjectRef] = None,
+    ) -> bool:
+        """Remove a tracked actor.
+
+        If the actor has already been started, this will stop the actor. This will
+        trigger the :meth:`TrackedActor.on_stop
+        <ray.air.execution._internal.tracked_actor.TrackedActor.on_stop>`
+        callback once the actor stopped.
+
+        If the actor has only been requested, but not started, yet, this will cancel
+        the actor request. This will not trigger any callback.
+
+        If ``kill=True``, this will use ``ray.kill()`` to forcefully terminate the
+        actor. Otherwise, graceful actor deconstruction will be scheduled after
+        all currently tracked futures are resolved.
+
+        This method returns a boolean, indicating if a stop future is tracked and
+        the ``on_stop`` callback will be invoked. If the actor has been alive,
+        this will be ``True``. If the actor hasn't been scheduled, yet, or failed
+        (and triggered the ``on_error`` callback), this will be ``False``.
+
+        Args:
+            tracked_actor: Tracked actor to be removed.
+            kill: If set, will forcefully terminate the actor instead of gracefully
+                scheduling termination.
+            stop_future: If set, use this future to track actor termination.
+                Otherwise, schedule a ``__ray_terminate__`` future.
+
+        Returns:
+            Boolean indicating if the actor was previously alive, and thus whether
+            a callback will be invoked once it is terminated.
+
+        """
+        if tracked_actor.actor_id in self._failed_actor_ids:
+            logger.debug(
+                f"Tracked actor already failed, no need to remove: {tracked_actor}"
+            )
+            return False
+        elif tracked_actor in self._live_actors_to_ray_actors_resources:
+            # Ray actor is running.
+
+            if not kill:
+                # Schedule __ray_terminate__ future
+                ray_actor, _ = self._live_actors_to_ray_actors_resources[tracked_actor]
+
+                # Clear state futures here to avoid resolving __ray_ready__ futures
+                for future in list(
+                    self._tracked_actors_to_state_futures[tracked_actor]
+                ):
+                    self._actor_state_events.discard_future(future)
+                    self._tracked_actors_to_state_futures[tracked_actor].remove(future)
+
+                    # If the __ray_ready__ future hasn't resolved yet, but we already
+                    # scheduled the actor via Actor.remote(), we just want to stop
+                    # it but not trigger any callbacks. This is in accordance with
+                    # the contract defined in the docstring.
+                    tracked_actor._on_start = None
+                    tracked_actor._on_stop = None
+                    tracked_actor._on_error = None
+
+                def on_actor_stop(*args, **kwargs):
+                    self._actor_stop_resolved(tracked_actor=tracked_actor)
+
+                if stop_future:
+                    # If the stop future was schedule via the actor manager,
+                    # discard (track it as state future instead).
+                    self._actor_task_events.discard_future(stop_future)
+                else:
+                    stop_future = ray_actor.__ray_terminate__.remote()
+
+                self._actor_state_events.track_future(
+                    future=stop_future,
+                    on_result=on_actor_stop,
+                    on_error=on_actor_stop,
+                )
+
+                self._tracked_actors_to_state_futures[tracked_actor].add(stop_future)
+            else:
+                # kill = True
+                self._live_actors_to_kill.add(tracked_actor)
+
+            return True
+
+        elif tracked_actor in self._pending_actors_to_attrs:
+            # Actor is pending, stop
+            _, _, resource_request = self._pending_actors_to_attrs.pop(tracked_actor)
+            self._resource_request_to_pending_actors[resource_request].remove(
+                tracked_actor
+            )
+            self._resource_manager.cancel_resource_request(
+                resource_request=resource_request
+            )
+            return False
+        else:
+            raise ValueError(f"Unknown tracked actor: {tracked_actor}")
+
+    def is_actor_started(self, tracked_actor: TrackedActor) -> bool:
+        """Returns True if the actor has been started.
+
+        Args:
+            tracked_actor: Tracked actor object.
+
+        Returns:
+            True if the actor has been started, False otherwise.
+        """
+        return (
+            tracked_actor in self._live_actors_to_ray_actors_resources
+            and tracked_actor.actor_id not in self._failed_actor_ids
+        )
+
+    def is_actor_failed(self, tracked_actor: TrackedActor) -> bool:
+        return tracked_actor.actor_id in self._failed_actor_ids
+
+    def get_actor_resources(
+        self, tracked_actor: TrackedActor
+    ) -> Optional[AcquiredResources]:
+        """Returns the acquired resources of an actor that has been started.
+
+        This will return ``None`` if the actor has not been started, yet.
+
+        Args:
+            tracked_actor: Tracked actor object.
+
+        Returns:
+            The acquired resources of the actor, or ``None`` if the actor has not
+            been started yet.
+        """
+        if not self.is_actor_started(tracked_actor):
+            return None
+
+        return self._live_actors_to_ray_actors_resources[tracked_actor][1]
+
+    def schedule_actor_task(
+        self,
+        tracked_actor: TrackedActor,
+        method_name: str,
+        args: Optional[Tuple] = None,
+        kwargs: Optional[Dict] = None,
+        on_result: Optional[Callable[[TrackedActor, Any], None]] = None,
+        on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
+        _return_future: bool = False,
+    ) -> Optional[ray.ObjectRef]:
+        """Schedule and track a task on an actor.
+
+        This method will schedule a remote task ``method_name`` on the
+        ``tracked_actor``.
+
+        This method accepts two optional callbacks that will be invoked when
+        their respective events are triggered.
+
+        The ``on_result`` callback is triggered when a task resolves successfully.
+        It should accept two arguments: The actor for which the
+        task resolved, and the result received from the remote call.
+
+        The ``on_error`` callback is triggered when a task fails.
+        It should accept two arguments: The actor for which the
+        task threw an error, and the exception.
+
+        Args:
+            tracked_actor: Actor to schedule task on.
+            method_name: Remote method name to invoke on the actor. If this is
+                e.g. ``foo``, then ``actor.foo.remote(*args, **kwargs)`` will be
+                scheduled.
+            args: Arguments to pass to the task.
+            kwargs: Keyword arguments to pass to the task.
+            on_result: Callback to invoke when the task resolves.
+            on_error: Callback to invoke when the task fails.
+            _return_future: If True, return the scheduled task's ``ObjectRef`` for
+                advanced callers. Defaults to False.
+
+        Raises:
+            ValueError: If the ``tracked_actor`` is not managed by this event manager.
+
+        Returns:
+            The scheduled task's ``ObjectRef`` if ``_return_future`` is True,
+            otherwise ``None``.
+        """
+        args = args or tuple()
+        kwargs = kwargs or {}
+
+        if tracked_actor.actor_id in self._failed_actor_ids:
+            return
+
+        tracked_actor_task = TrackedActorTask(
+            tracked_actor=tracked_actor, on_result=on_result, on_error=on_error
+        )
+
+        if tracked_actor not in self._live_actors_to_ray_actors_resources:
+            # Actor is not started, yet
+            if tracked_actor not in self._pending_actors_to_attrs:
+                raise ValueError(
+                    f"Tracked actor is not managed by this event manager: "
+                    f"{tracked_actor}"
+                )
+
+            # Cache tasks for future execution
+            self._pending_actors_to_enqueued_actor_tasks[tracked_actor].append(
+                (tracked_actor_task, method_name, args, kwargs)
+            )
+        else:
+            res = self._schedule_tracked_actor_task(
+                tracked_actor_task=tracked_actor_task,
+                method_name=method_name,
+                args=args,
+                kwargs=kwargs,
+                _return_future=_return_future,
+            )
+            if _return_future:
+                return res[1]
+
+    def _schedule_tracked_actor_task(
+        self,
+        tracked_actor_task: TrackedActorTask,
+        method_name: str,
+        *,
+        args: Optional[Tuple] = None,
+        kwargs: Optional[Dict] = None,
+        _return_future: bool = False,
+    ) -> Union[TrackedActorTask, Tuple[TrackedActorTask, ray.ObjectRef]]:
+        tracked_actor = tracked_actor_task._tracked_actor
+        ray_actor, _ = self._live_actors_to_ray_actors_resources[tracked_actor]
+
+        try:
+            remote_fn = getattr(ray_actor, method_name)
+        except AttributeError as e:
+            raise AttributeError(
+                f"Remote function `{method_name}()` does not exist for this actor."
+            ) from e
+
+        def on_result(result: Any):
+            self._actor_task_resolved(
+                tracked_actor_task=tracked_actor_task, result=result
+            )
+
+        def on_error(exception: Exception):
+            self._actor_task_failed(
+                tracked_actor_task=tracked_actor_task, exception=exception
+            )
+
+        future = remote_fn.remote(*args, **kwargs)
+
+        self._actor_task_events.track_future(
+            future=future, on_result=on_result, on_error=on_error
+        )
+
+        self._tracked_actors_to_task_futures[tracked_actor].add(future)
+
+        if _return_future:
+            return tracked_actor_task, future
+
+        return tracked_actor_task
+
+    def schedule_actor_tasks(
+        self,
+        tracked_actors: List[TrackedActor],
+        method_name: str,
+        *,
+        args: Optional[Union[Tuple, List[Tuple]]] = None,
+        kwargs: Optional[Union[Dict, List[Dict]]] = None,
+        on_result: Optional[Callable[[TrackedActor, Any], None]] = None,
+        on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
+    ) -> None:
+        """Schedule and track tasks on a list of actors.
+
+        This method will schedule a remote task ``method_name`` on all
+        ``tracked_actors``.
+
+        ``args`` and ``kwargs`` can be a single tuple/dict, in which case the same
+        (keyword) arguments are passed to all actors. If a list is passed instead,
+        they are mapped to the respective actors. In that case, the list of
+        (keyword) arguments must be the same length as the list of actors.
+
+        This method accepts two optional callbacks that will be invoked when
+        their respective events are triggered.
+
+        The ``on_result`` callback is triggered when a task resolves successfully.
+        It should accept two arguments: The actor for which the
+        task resolved, and the result received from the remote call.
+
+        The ``on_error`` callback is triggered when a task fails.
+        It should accept two arguments: The actor for which the
+        task threw an error, and the exception.
+
+        Args:
+            tracked_actors: List of actors to schedule tasks on.
+            method_name: Remote actor method to invoke on the actors. If this is
+                e.g. ``foo``, then ``actor.foo.remote(*args, **kwargs)`` will be
+                scheduled on all actors.
+            args: Arguments to pass to the task.
+            kwargs: Keyword arguments to pass to the task.
+            on_result: Callback to invoke when the task resolves.
+            on_error: Callback to invoke when the task fails.
+
+        """
+        if not isinstance(args, List):
+            args_list = [args] * len(tracked_actors)
+        else:
+            if len(tracked_actors) != len(args):
+                raise ValueError(
+                    f"Length of args must be the same as tracked_actors "
+                    f"list. Got `len(kwargs)={len(kwargs)}` and "
+                    f"`len(tracked_actors)={len(tracked_actors)}"
+                )
+            args_list = args
+
+        if not isinstance(kwargs, List):
+            kwargs_list = [kwargs] * len(tracked_actors)
+        else:
+            if len(tracked_actors) != len(kwargs):
+                raise ValueError(
+                    f"Length of kwargs must be the same as tracked_actors "
+                    f"list. Got `len(args)={len(args)}` and "
+                    f"`len(tracked_actors)={len(tracked_actors)}"
+                )
+            kwargs_list = kwargs
+
+        for tracked_actor, args, kwargs in zip(tracked_actors, args_list, kwargs_list):
+            self.schedule_actor_task(
+                tracked_actor=tracked_actor,
+                method_name=method_name,
+                args=args,
+                kwargs=kwargs,
+                on_result=on_result,
+                on_error=on_error,
+            )
+
+    def clear_actor_task_futures(self, tracked_actor: TrackedActor):
+        """Discard all actor task futures from a tracked actor."""
+        futures = self._tracked_actors_to_task_futures.pop(tracked_actor, [])
+        for future in futures:
+            self._actor_task_events.discard_future(future)
+
+    def _cleanup_actor_futures(self, tracked_actor: TrackedActor):
+        # Remove all actor task futures
+        self.clear_actor_task_futures(tracked_actor=tracked_actor)
+
+        # Remove all actor state futures
+        futures = self._tracked_actors_to_state_futures.pop(tracked_actor, [])
+        for future in futures:
+            self._actor_state_events.discard_future(future)
+
+    def cleanup(self):
+        for (
+            actor,
+            acquired_resources,
+        ) in self._live_actors_to_ray_actors_resources.values():
+            ray.kill(actor)
+            self._resource_manager.free_resources(acquired_resources)
+
+        for (
+            resource_request,
+            pending_actors,
+        ) in self._resource_request_to_pending_actors.items():
+            for i in range(len(pending_actors)):
+                self._resource_manager.cancel_resource_request(resource_request)
+
+        self._resource_manager.clear()
+
+        self.__init__(resource_manager=self._resource_manager)

python/ray/air/execution/_internal/barrier.py

@@ -0,0 +1,93 @@
+from typing import Any, Callable, List, Optional, Tuple
+
+
+class Barrier:
+    """Barrier to collect results and process them in bulk.
+
+    A barrier can be used to collect multiple results and process them in bulk once
+    a certain count or a timeout is reached.
+
+    For instance, if ``max_results=N``, the ``on_completion`` callback will be
+    invoked once :meth:`arrive` has been called ``N`` times.
+
+    The completion callback will only be invoked once, even if more results
+    arrive after completion. The collected results can be resetted
+    with :meth:`reset`, after which the callback may be invoked again.
+
+    The completion callback should expect one argument, which is the barrier
+    object that completed.
+
+    Args:
+        max_results: Maximum number of results to collect before a call to
+            :meth:`wait` resolves or the :meth:`on_completion` callback is invoked.
+        on_completion: Callback to invoke when ``max_results`` results
+            arrived at the barrier.
+
+    """
+
+    def __init__(
+        self,
+        max_results: int,
+        *,
+        on_completion: Optional[Callable[["Barrier"], None]] = None,
+    ):
+        self._max_results = max_results
+
+        # on_completion callback
+        self._completed = False
+        self._on_completion = on_completion
+
+        # Collect received results
+        self._results: List[Tuple[Any]] = []
+
+    def arrive(self, *data: Any):
+        """Notify barrier that a result successfully arrived.
+
+        This will count against the ``max_results`` limit. The received result
+        will be included in a call to :meth:`get_results`.
+
+        Args:
+            *data: Result data to be cached. Can be obtained via :meth:`get_results`.
+
+        """
+        if len(data) == 1:
+            data = data[0]
+
+        self._results.append(data)
+        self._check_completion()
+
+    def _check_completion(self):
+        if self._completed:
+            # Already fired completion callback
+            return
+
+        if self.num_results >= self._max_results:
+            # Barrier is complete
+            self._completed = True
+
+            if self._on_completion:
+                self._on_completion(self)
+
+    @property
+    def completed(self) -> bool:
+        """Returns True if the barrier is completed."""
+        return self._completed
+
+    @property
+    def num_results(self) -> int:
+        """Number of received (successful) results."""
+        return len(self._results)
+
+    def get_results(self) -> List[Tuple[Any]]:
+        """Return list of received results."""
+        return self._results
+
+    def reset(self) -> None:
+        """Reset barrier, removing all received results.
+
+        Resetting the barrier will reset the completion status. When ``max_results``
+        is set and enough new events arrive after resetting, the
+        :meth:`on_completion` callback will be invoked again.
+        """
+        self._completed = False
+        self._results = []

python/ray/air/execution/_internal/event_manager.py

@@ -0,0 +1,148 @@
+import random
+from typing import Any, Callable, Dict, Iterable, Optional, Set, Tuple, Union
+
+import ray
+
+_ResultCallback = Callable[[Any], None]
+_ErrorCallback = Callable[[Exception], None]
+
+
+class RayEventManager:
+    """Event manager for Ray futures.
+
+    The event manager can be used to track futures and invoke callbacks when
+    they resolve.
+
+    Futures are tracked with :meth:`track_future`. Future can then be awaited with
+    :meth:`wait`. When futures successfully resolve, they trigger an optional
+    ``on_result`` callback that can be passed to :meth:`track_future`. If they
+    fail, they trigger an optional ``on_error`` callback.
+
+    Args:
+        shuffle_futures: If True, futures will be shuffled before awaited. This
+            will avoid implicit prioritization of futures within Ray.
+    """
+
+    def __init__(self, shuffle_futures: bool = True):
+        self._shuffle_futures = shuffle_futures
+
+        # Map of futures to callbacks (result, error)
+        self._tracked_futures: Dict[
+            ray.ObjectRef, Tuple[Optional[_ResultCallback], Optional[_ErrorCallback]]
+        ] = {}
+
+    def track_future(
+        self,
+        future: ray.ObjectRef,
+        on_result: Optional[_ResultCallback] = None,
+        on_error: Optional[_ErrorCallback] = None,
+    ):
+        """Track a single future and invoke callbacks on resolution.
+
+        Control has to be yielded to the event manager for the callbacks to
+        be invoked, either via :meth:`wait` or via :meth:`resolve_future`.
+
+        Args:
+            future: Ray future to await.
+            on_result: Callback to invoke when the future resolves successfully.
+            on_error: Callback to invoke when the future fails.
+
+        """
+        self._tracked_futures[future] = (on_result, on_error)
+
+    def track_futures(
+        self,
+        futures: Iterable[ray.ObjectRef],
+        on_result: Optional[_ResultCallback] = None,
+        on_error: Optional[_ErrorCallback] = None,
+    ):
+        """Track multiple futures and invoke callbacks on resolution.
+
+        Control has to be yielded to the event manager for the callbacks to
+        be invoked, either via :meth:`wait` or via :meth:`resolve_future`.
+
+        Args:
+            futures: Ray futures to await.
+            on_result: Callback to invoke when the future resolves successfully.
+            on_error: Callback to invoke when the future fails.
+
+        """
+        for future in futures:
+            self.track_future(future, on_result=on_result, on_error=on_error)
+
+    def discard_future(self, future: ray.ObjectRef):
+        """Remove future from tracking.
+
+        The future will not be awaited anymore, and it will not trigger any callbacks.
+
+        Args:
+            future: Ray futures to discard.
+        """
+        self._tracked_futures.pop(future, None)
+
+    def get_futures(self) -> Set[ray.ObjectRef]:
+        """Get futures tracked by the event manager."""
+        return set(self._tracked_futures)
+
+    @property
+    def num_futures(self) -> int:
+        return len(self._tracked_futures)
+
+    def resolve_future(self, future: ray.ObjectRef):
+        """Resolve a single future.
+
+        This method will block until the future is available. It will then
+        trigger the callback associated to the future and the event (success
+        or error), if specified.
+
+        Args:
+            future: Ray future to resolve.
+
+        """
+        try:
+            on_result, on_error = self._tracked_futures.pop(future)
+        except KeyError as e:
+            raise ValueError(
+                f"Future {future} is not tracked by this RayEventManager"
+            ) from e
+
+        try:
+            result = ray.get(future)
+        except Exception as e:
+            if on_error:
+                on_error(e)
+            else:
+                raise e
+        else:
+            if on_result:
+                on_result(result)
+
+    def wait(
+        self,
+        timeout: Optional[Union[float, int]] = None,
+        num_results: Optional[int] = 1,
+    ):
+        """Wait up to ``timeout`` seconds for ``num_results`` futures to resolve.
+
+        If ``timeout=None``, this method will block until all `num_results`` futures
+        resolve. If ``num_results=None``, this method will await all tracked futures.
+
+        For every future that resolves, the respective associated callbacks will be
+        invoked.
+
+        Args:
+            timeout: Timeout in second to wait for futures to resolve.
+            num_results: Number of futures to await. If ``None``, will wait for
+                all tracked futures to resolve.
+
+        """
+        futures = list(self.get_futures())
+
+        if self._shuffle_futures:
+            random.shuffle(futures)
+
+        num_results = num_results or len(futures)
+
+        ready, _ = ray.wait(list(futures), timeout=timeout, num_returns=num_results)
+        for future in ready:
+            self.resolve_future(future)

python/ray/air/execution/_internal/tracked_actor.py

@@ -0,0 +1,62 @@
+from typing import Callable, Optional
+
+
+class TrackedActor:
+    """Actor tracked by an actor manager.
+
+    This object is used to reference a Ray actor on an actor manager
+
+    Existence of this object does not mean that the Ray actor has already been started.
+    Actor state can be inquired from the actor manager tracking the Ray actor.
+
+    Note:
+        Objects of this class are returned by the :class:`RayActorManager`.
+        This class should not be instantiated manually.
+
+    Attributes:
+        actor_id: ID for identification of the actor within the actor manager. This
+            ID is not related to the Ray actor ID.
+
+    """
+
+    def __init__(
+        self,
+        actor_id: int,
+        on_start: Optional[Callable[["TrackedActor"], None]] = None,
+        on_stop: Optional[Callable[["TrackedActor"], None]] = None,
+        on_error: Optional[Callable[["TrackedActor", Exception], None]] = None,
+    ):
+        """Initialize the tracked actor.
+
+        Args:
+            actor_id: ID for identification of the actor within the actor manager.
+            on_start: Callback to invoke when the actor started.
+            on_stop: Callback to invoke when the actor stopped.
+            on_error: Callback to invoke when the actor failed.
+        """
+        self.actor_id = actor_id
+        self._on_start = on_start
+        self._on_stop = on_stop
+        self._on_error = on_error
+
+    def set_on_start(self, on_start: Optional[Callable[["TrackedActor"], None]]):
+        self._on_start = on_start
+
+    def set_on_stop(self, on_stop: Optional[Callable[["TrackedActor"], None]]):
+        self._on_stop = on_stop
+
+    def set_on_error(
+        self, on_error: Optional[Callable[["TrackedActor", Exception], None]]
+    ):
+        self._on_error = on_error
+
+    def __repr__(self):
+        return f"<TrackedActor {self.actor_id}>"
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        return self.actor_id == other.actor_id
+
+    def __hash__(self):
+        return hash(self.actor_id)

python/ray/air/execution/_internal/tracked_actor_task.py

@@ -0,0 +1,42 @@
+from typing import Any, Callable, Optional
+
+from ray.air.execution._internal.tracked_actor import TrackedActor
+
+
+class TrackedActorTask:
+    """Actor task tracked by a Ray event manager.
+
+    This container class is used to define callbacks to be invoked when
+    the task resolves, errors, or times out.
+
+    Note:
+        Objects of this class are returned by the :class:`RayActorManager`.
+        This class should not be instantiated manually.
+
+    Args:
+        tracked_actor: Tracked actor object this task is scheduled on.
+        on_result: Callback to invoke when the task resolves.
+        on_error: Callback to invoke when the task fails.
+
+    Example:
+
+        .. code-block:: python
+
+            tracked_futures = actor_manager.schedule_actor_tasks(
+                actor_manager.live_actors,
+                "foo",
+                on_result=lambda actor, result: print(result)
+                )
+
+    """
+
+    def __init__(
+        self,
+        tracked_actor: TrackedActor,
+        on_result: Optional[Callable[[TrackedActor, Any], None]] = None,
+        on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
+    ):
+        self._tracked_actor = tracked_actor
+
+        self._on_result = on_result
+        self._on_error = on_error

python/ray/air/execution/resources/__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",
+]

python/ray/air/execution/resources/fixed.py

@@ -0,0 +1,147 @@
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import ray
+from ray import SCRIPT_MODE
+from ray.air.execution.resources.request import (
+    AcquiredResources,
+    RemoteRayEntity,
+    ResourceRequest,
+)
+from ray.air.execution.resources.resource_manager import ResourceManager
+from ray.util.annotations import DeveloperAPI
+
+# Avoid numerical errors by multiplying and subtracting with this number.
+# Compare: 0.99 - 0.33 = 0.65999... vs (0.99 * 1000 - 0.33 * 1000) / 1000 = 0.66
+_DIGITS = 100000
+
+
+@DeveloperAPI
+@dataclass
+class FixedAcquiredResources(AcquiredResources):
+    bundles: List[Dict[str, float]]
+
+    def _annotate_remote_entity(
+        self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int
+    ) -> RemoteRayEntity:
+        bundle = bundle.copy()
+        num_cpus = bundle.pop("CPU", 0)
+        num_gpus = bundle.pop("GPU", 0)
+        memory = bundle.pop("memory", 0.0)
+
+        return entity.options(
+            num_cpus=num_cpus,
+            num_gpus=num_gpus,
+            memory=memory,
+            resources=bundle,
+        )
+
+
+@DeveloperAPI
+class FixedResourceManager(ResourceManager):
+    """Fixed budget based resource manager.
+
+    This resource manager keeps track of a fixed set of resources. When resources
+    are acquired, they are subtracted from the budget. When resources are freed,
+    they are added back to the budget.
+
+    The resource manager still requires resources to be requested before they become
+    available. However, because the resource requests are virtual, this will not
+    trigger autoscaling.
+
+    Additionally, resources are not reserved on request, only on acquisition. Thus,
+    acquiring a resource can change the availability of other requests. Note that
+    this behavior may be changed in future implementations.
+
+    The fixed resource manager does not support placement strategies. Using
+    ``STRICT_SPREAD`` will result in an error. ``STRICT_PACK`` will succeed only
+    within a placement group bundle. All other placement group arguments will be
+    ignored.
+
+    Args:
+        total_resources: Budget of resources to manage. Defaults to all available
+            resources in the current task or all cluster resources (if outside a task).
+
+    """
+
+    _resource_cls: AcquiredResources = FixedAcquiredResources
+
+    def __init__(self, total_resources: Optional[Dict[str, float]] = None):
+        rtc = ray.get_runtime_context()
+
+        if not total_resources:
+            if rtc.worker.mode in {None, SCRIPT_MODE}:
+                total_resources = ray.cluster_resources()
+            else:
+                total_resources = rtc.get_assigned_resources()
+
+        # If we are in a placement group, all of our resources will be in a bundle
+        # and thus fulfill requirements of STRICT_PACK - but only if child tasks
+        # are captured by the pg.
+        self._allow_strict_pack = (
+            ray.util.get_current_placement_group() is not None
+            and rtc.should_capture_child_tasks_in_placement_group
+        )
+
+        self._total_resources = total_resources
+        self._requested_resources = []
+        self._used_resources = []
+
+    @property
+    def _available_resources(self) -> Dict[str, float]:
+        available_resources = self._total_resources.copy()
+
+        for used_resources in self._used_resources:
+            all_resources = used_resources.required_resources
+            for k, v in all_resources.items():
+                available_resources[k] = (
+                    available_resources[k] * _DIGITS - v * _DIGITS
+                ) / _DIGITS
+        return available_resources
+
+    def request_resources(self, resource_request: ResourceRequest):
+        if resource_request.strategy == "STRICT_SPREAD" or (
+            not self._allow_strict_pack and resource_request.strategy == "STRICT_PACK"
+        ):
+            raise RuntimeError(
+                f"Requested a resource with placement strategy "
+                f"{resource_request.strategy}, but this cannot be fulfilled by a "
+                f"FixedResourceManager. In a nested setting, please set the inner "
+                f"placement strategy to be less restrictive (i.e. no STRICT_ strategy)."
+            )
+
+        self._requested_resources.append(resource_request)
+
+    def cancel_resource_request(self, resource_request: ResourceRequest):
+        self._requested_resources.remove(resource_request)
+
+    def has_resources_ready(self, resource_request: ResourceRequest) -> bool:
+        if resource_request not in self._requested_resources:
+            return False
+
+        available_resources = self._available_resources
+        all_resources = resource_request.required_resources
+        for k, v in all_resources.items():
+            if available_resources.get(k, 0.0) < v:
+                return False
+        return True
+
+    def acquire_resources(
+        self, resource_request: ResourceRequest
+    ) -> Optional[AcquiredResources]:
+        if not self.has_resources_ready(resource_request):
+            return None
+
+        self._used_resources.append(resource_request)
+        return self._resource_cls(
+            bundles=resource_request.bundles, resource_request=resource_request
+        )
+
+    def free_resources(self, acquired_resource: AcquiredResources):
+        resources = acquired_resource.resource_request
+        self._used_resources.remove(resources)
+
+    def clear(self):
+        # Reset internal state
+        self._requested_resources = []
+        self._used_resources = []

python/ray/air/execution/resources/placement_group.py

@@ -0,0 +1,214 @@
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Set
+
+import ray
+from ray.air.execution.resources.request import (
+    AcquiredResources,
+    RemoteRayEntity,
+    ResourceRequest,
+)
+from ray.air.execution.resources.resource_manager import ResourceManager
+from ray.util.annotations import DeveloperAPI
+from ray.util.placement_group import PlacementGroup, remove_placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+
+@DeveloperAPI
+@dataclass
+class PlacementGroupAcquiredResources(AcquiredResources):
+    placement_group: PlacementGroup
+
+    def _annotate_remote_entity(
+        self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int
+    ) -> RemoteRayEntity:
+        bundle = bundle.copy()
+        num_cpus = bundle.pop("CPU", 0)
+        num_gpus = bundle.pop("GPU", 0)
+        memory = bundle.pop("memory", 0.0)
+
+        return entity.options(
+            scheduling_strategy=PlacementGroupSchedulingStrategy(
+                placement_group=self.placement_group,
+                placement_group_bundle_index=bundle_index,
+                placement_group_capture_child_tasks=True,
+            ),
+            num_cpus=num_cpus,
+            num_gpus=num_gpus,
+            memory=memory,
+            resources=bundle,
+        )
+
+
+@DeveloperAPI
+class PlacementGroupResourceManager(ResourceManager):
+    """Resource manager using placement groups as the resource backend.
+
+    This manager will use placement groups to fulfill resource requests. Requesting
+    a resource will schedule the placement group. Acquiring a resource will
+    return a ``PlacementGroupAcquiredResources`` that can be used to schedule
+    Ray tasks and actors on the placement group. Freeing an acquired resource
+    will destroy the associated placement group.
+
+    Ray core does not emit events when resources are available. Instead, the
+    scheduling state has to be periodically updated.
+
+    Per default, placement group scheduling state is refreshed every time when
+    resource state is inquired, but not more often than once every ``update_interval_s``
+    seconds. Alternatively, staging futures can be retrieved (and awaited) with
+    ``get_resource_futures()`` and state update can be force with ``update_state()``.
+
+    Args:
+        update_interval_s: Minimum interval in seconds between updating scheduling
+            state of placement groups.
+
+    """
+
+    _resource_cls: AcquiredResources = PlacementGroupAcquiredResources
+
+    def __init__(self, update_interval_s: float = 0.1):
+        # Internally, the placement group lifecycle is like this:
+        # - Resources are requested with ``request_resources()``
+        # - A placement group is scheduled ("staged")
+        # - A ``PlacementGroup.ready()`` future is scheduled ("staging future")
+        # - We update the scheduling state when we need to
+        #   (e.g. when ``has_resources_ready()`` is called)
+        # - When staging futures resolve, a placement group is moved from "staging"
+        #   to "ready"
+        # - When a resource request is canceled, we remove a placement group from
+        #   "staging". If there are not staged placement groups
+        #   (because they are already "ready"), we remove one from "ready" instead.
+        # - When a resource is acquired, the pg is removed from "ready" and moved
+        #   to "acquired"
+        # - When a resource is freed, the pg is removed from "acquired" and destroyed
+
+        # Mapping of placement group to request
+        self._pg_to_request: Dict[PlacementGroup, ResourceRequest] = {}
+
+        # PGs that are staged but not "ready", yet (i.e. not CREATED)
+        self._request_to_staged_pgs: Dict[
+            ResourceRequest, Set[PlacementGroup]
+        ] = defaultdict(set)
+
+        # PGs that are CREATED and can be used by tasks and actors
+        self._request_to_ready_pgs: Dict[
+            ResourceRequest, Set[PlacementGroup]
+        ] = defaultdict(set)
+
+        # Staging futures used to update internal state.
+        # We keep a double mapping here for better lookup efficiency.
+        self._staging_future_to_pg: Dict[ray.ObjectRef, PlacementGroup] = dict()
+        self._pg_to_staging_future: Dict[PlacementGroup, ray.ObjectRef] = dict()
+
+        # Set of acquired PGs. We keep track of these here to make sure we
+        # only free PGs that this manager managed.
+        self._acquired_pgs: Set[PlacementGroup] = set()
+
+        # Minimum time between updates of the internal state
+        self.update_interval_s = update_interval_s
+        self._last_update = time.monotonic() - self.update_interval_s - 1
+
+    def get_resource_futures(self) -> List[ray.ObjectRef]:
+        return list(self._staging_future_to_pg.keys())
+
+    def _maybe_update_state(self):
+        now = time.monotonic()
+        if now > self._last_update + self.update_interval_s:
+            self.update_state()
+
+    def update_state(self):
+        ready, not_ready = ray.wait(
+            list(self._staging_future_to_pg.keys()),
+            num_returns=len(self._staging_future_to_pg),
+            timeout=0,
+        )
+        for future in ready:
+            # Remove staging future
+            pg = self._staging_future_to_pg.pop(future)
+            self._pg_to_staging_future.pop(pg)
+            # Fetch resource request
+            request = self._pg_to_request[pg]
+            # Remove from staging, add to ready
+            self._request_to_staged_pgs[request].remove(pg)
+            self._request_to_ready_pgs[request].add(pg)
+        self._last_update = time.monotonic()
+
+    def request_resources(self, resource_request: ResourceRequest):
+        pg = resource_request.to_placement_group()
+        self._pg_to_request[pg] = resource_request
+        self._request_to_staged_pgs[resource_request].add(pg)
+
+        future = pg.ready()
+        self._staging_future_to_pg[future] = pg
+        self._pg_to_staging_future[pg] = future
+
+    def cancel_resource_request(self, resource_request: ResourceRequest):
+        if self._request_to_staged_pgs[resource_request]:
+            pg = self._request_to_staged_pgs[resource_request].pop()
+
+            # PG was staging
+            future = self._pg_to_staging_future.pop(pg)
+            self._staging_future_to_pg.pop(future)
+
+            # Cancel the pg.ready task.
+            # Otherwise, it will be pending node assignment forever.
+            ray.cancel(future)
+        else:
+            # PG might be ready
+            pg = self._request_to_ready_pgs[resource_request].pop()
+            if not pg:
+                raise RuntimeError(
+                    "Cannot cancel resource request: No placement group was "
+                    f"staged or is ready. Make sure to not cancel more resource "
+                    f"requests than you've created. Request: {resource_request}"
+                )
+
+        self._pg_to_request.pop(pg)
+        ray.util.remove_placement_group(pg)
+
+    def has_resources_ready(self, resource_request: ResourceRequest) -> bool:
+        if not bool(len(self._request_to_ready_pgs[resource_request])):
+            # Only update state if needed
+            self._maybe_update_state()
+
+        return bool(len(self._request_to_ready_pgs[resource_request]))
+
+    def acquire_resources(
+        self, resource_request: ResourceRequest
+    ) -> Optional[PlacementGroupAcquiredResources]:
+        if not self.has_resources_ready(resource_request):
+            return None
+
+        pg = self._request_to_ready_pgs[resource_request].pop()
+        self._acquired_pgs.add(pg)
+
+        return self._resource_cls(placement_group=pg, resource_request=resource_request)
+
+    def free_resources(self, acquired_resource: PlacementGroupAcquiredResources):
+        pg = acquired_resource.placement_group
+
+        self._acquired_pgs.remove(pg)
+        remove_placement_group(pg)
+        self._pg_to_request.pop(pg)
+
+    def clear(self):
+        if not ray.is_initialized():
+            return
+
+        for staged_pgs in self._request_to_staged_pgs.values():
+            for staged_pg in staged_pgs:
+                remove_placement_group(staged_pg)
+
+        for ready_pgs in self._request_to_ready_pgs.values():
+            for ready_pg in ready_pgs:
+                remove_placement_group(ready_pg)
+
+        for acquired_pg in self._acquired_pgs:
+            remove_placement_group(acquired_pg)
+
+        # Reset internal state
+        self.__init__(update_interval_s=self.update_interval_s)
+
+    def __del__(self):
+        self.clear()

python/ray/air/execution/resources/request.py

@@ -0,0 +1,259 @@
+import abc
+import json
+from copy import deepcopy
+from dataclasses import dataclass
+from inspect import signature
+from typing import Dict, List, Union
+
+import ray
+from ray.util import placement_group
+from ray.util.annotations import DeveloperAPI
+
+RemoteRayEntity = Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass]
+
+
+def _sum_bundles(bundles: List[Dict[str, float]]) -> Dict[str, float]:
+    """Sum all resources in a list of resource bundles.
+
+    Args:
+        bundles: List of resource bundles.
+
+    Returns:
+        Dict containing all resources summed up.
+    """
+    resources = {}
+    for bundle in bundles:
+        for k, v in bundle.items():
+            resources[k] = resources.get(k, 0) + v
+    return resources
+
+
+@DeveloperAPI
+class ResourceRequest:
+    """Request for resources.
+
+    This class is used to define a resource request. A resource request comprises one
+    or more bundles of resources and instructions on the scheduling behavior.
+
+    The resource request can be submitted to a resource manager, which will
+    schedule the resources. Depending on the resource backend, this may instruct
+    Ray to scale up (autoscaling).
+
+    Resource requests are compatible with the most fine-grained low-level resource
+    backend, which are Ray placement groups.
+
+    Args:
+        bundles: A list of bundles which represent the resources requirements.
+            E.g. ``[{"CPU": 1, "GPU": 1}]``.
+        strategy: The scheduling strategy to acquire the bundles.
+
+         - "PACK": Packs Bundles into as few nodes as possible.
+         - "SPREAD": Places Bundles across distinct nodes as even as possible.
+         - "STRICT_PACK": Packs Bundles into one node. The group is
+           not allowed to span multiple nodes.
+         - "STRICT_SPREAD": Packs Bundles across distinct nodes.
+        *args: Passed to the call of ``placement_group()``, if applicable.
+        **kwargs: Passed to the call of ``placement_group()``, if applicable.
+
+    """
+
+    def __init__(
+        self,
+        bundles: List[Dict[str, Union[int, float]]],
+        strategy: str = "PACK",
+        *args,
+        **kwargs,
+    ):
+        if not bundles:
+            raise ValueError("Cannot initialize a ResourceRequest with zero bundles.")
+
+        # Remove empty resource keys
+        self._bundles = [
+            {k: float(v) for k, v in bundle.items() if v != 0} for bundle in bundles
+        ]
+
+        # Check if the head bundle is empty (no resources defined or all resources
+        # are 0 (and thus removed in the previous step)
+        if not self._bundles[0]:
+            # This is when the head bundle doesn't need resources.
+            self._head_bundle_is_empty = True
+            self._bundles.pop(0)
+
+            if not self._bundles:
+                raise ValueError(
+                    "Cannot initialize a ResourceRequest with an empty head "
+                    "and zero worker bundles."
+                )
+        else:
+            self._head_bundle_is_empty = False
+
+        self._strategy = strategy
+        self._args = args
+        self._kwargs = kwargs
+
+        self._hash = None
+        self._bound = None
+
+        self._bind()
+
+    @property
+    def head_bundle_is_empty(self):
+        """Returns True if head bundle is empty while child bundles
+        need resources.
+
+        This is considered an internal API within Tune.
+        """
+        return self._head_bundle_is_empty
+
+    @property
+    @DeveloperAPI
+    def head_cpus(self) -> float:
+        """Returns the number of cpus in the head bundle."""
+        return 0.0 if self._head_bundle_is_empty else self._bundles[0].get("CPU", 0.0)
+
+    @property
+    @DeveloperAPI
+    def bundles(self) -> List[Dict[str, float]]:
+        """Returns a deep copy of resource bundles"""
+        return deepcopy(self._bundles)
+
+    @property
+    def required_resources(self) -> Dict[str, float]:
+        """Returns a dict containing the sums of all resources"""
+        return _sum_bundles(self._bundles)
+
+    @property
+    @DeveloperAPI
+    def strategy(self) -> str:
+        """Returns the placement strategy"""
+        return self._strategy
+
+    def _bind(self):
+        """Bind the args and kwargs to the `placement_group()` signature.
+
+        We bind the args and kwargs, so we can compare equality of two resource
+        requests. The main reason for this is that the `placement_group()` API
+        can evolve independently from the ResourceRequest API (e.g. adding new
+        arguments). Then, `ResourceRequest(bundles, strategy, arg=arg)` should
+        be the same as `ResourceRequest(bundles, strategy, arg)`.
+        """
+        sig = signature(placement_group)
+        try:
+            self._bound = sig.bind(
+                self._bundles, self._strategy, *self._args, **self._kwargs
+            )
+        except Exception as exc:
+            raise RuntimeError(
+                "Invalid definition for resource request. Please check "
+                "that you passed valid arguments to the ResourceRequest "
+                "object."
+            ) from exc
+
+    def to_placement_group(self):
+        return placement_group(*self._bound.args, **self._bound.kwargs)
+
+    def __eq__(self, other: "ResourceRequest"):
+        return (
+            isinstance(other, ResourceRequest)
+            and self._bound == other._bound
+            and self.head_bundle_is_empty == other.head_bundle_is_empty
+        )
+
+    def __hash__(self):
+        if not self._hash:
+            # Cache hash
+            self._hash = hash(
+                json.dumps(
+                    {"args": self._bound.args, "kwargs": self._bound.kwargs},
+                    sort_keys=True,
+                    indent=0,
+                    ensure_ascii=True,
+                )
+            )
+        return self._hash
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop("_hash", None)
+        state.pop("_bound", None)
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._hash = None
+        self._bound = None
+        self._bind()
+
+    def __repr__(self) -> str:
+        return (
+            f"<ResourceRequest (_bound={self._bound}, "
+            f"head_bundle_is_empty={self.head_bundle_is_empty})>"
+        )
+
+
+@DeveloperAPI
+@dataclass
+class AcquiredResources(abc.ABC):
+    """Base class for resources that have been acquired.
+
+    Acquired resources can be associated to Ray objects, which can then be
+    scheduled using these resources.
+
+    Internally this can point e.g. to a placement group, a placement
+    group bundle index, or just raw resources.
+
+    The main API is the `annotate_remote_entities` method. This will associate
+    remote Ray objects (tasks and actors) with the acquired resources by setting
+    the Ray remote options to use the acquired resources.
+    """
+
+    resource_request: ResourceRequest
+
+    def annotate_remote_entities(
+        self, entities: List[RemoteRayEntity]
+    ) -> List[Union[RemoteRayEntity]]:
+        """Return remote ray entities (tasks/actors) to use the acquired resources.
+
+        The first entity will be associated with the first bundle, the second
+        entity will be associated with the second bundle, etc.
+
+        Args:
+            entities: Remote Ray entities to annotate with the acquired resources.
+
+        Returns:
+            The list of annotated remote Ray entities.
+        """
+        bundles = self.resource_request.bundles
+
+        # Also count the empty head bundle as a bundle
+        num_bundles = len(bundles) + int(self.resource_request.head_bundle_is_empty)
+
+        if len(entities) > num_bundles:
+            raise RuntimeError(
+                f"The number of callables to annotate ({len(entities)}) cannot "
+                f"exceed the number of available bundles ({num_bundles})."
+            )
+
+        annotated = []
+
+        if self.resource_request.head_bundle_is_empty:
+            # The empty head bundle is place on the first bundle index with empty
+            # resources.
+            annotated.append(
+                self._annotate_remote_entity(entities[0], {}, bundle_index=0)
+            )
+
+            # Shift the remaining entities
+            entities = entities[1:]
+
+        for i, (entity, bundle) in enumerate(zip(entities, bundles)):
+            annotated.append(
+                self._annotate_remote_entity(entity, bundle, bundle_index=i)
+            )
+
+        return annotated
+
+    def _annotate_remote_entity(
+        self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int
+    ) -> RemoteRayEntity:
+        raise NotImplementedError

python/ray/air/execution/resources/resource_manager.py

@@ -0,0 +1,155 @@
+import abc
+from typing import List, Optional
+
+import ray
+from ray.air.execution.resources.request import AcquiredResources, ResourceRequest
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class ResourceManager(abc.ABC):
+    """Resource manager interface.
+
+    A resource manager can be used to request resources from a Ray cluster and
+    allocate them to remote Ray tasks or actors.
+
+    Resources have to be requested before they can be acquired.
+
+    Resources managed by the resource manager can be in three states:
+
+    1. "Requested":  The resources have been requested but are not yet available to
+       schedule remote Ray objects. The resource request may trigger autoscaling,
+       and can be cancelled if no longer needed.
+    2. "Ready": The requested resources are now available to schedule remote Ray
+       objects. They can be acquired and subsequently used remote Ray objects.
+       The resource request can still be cancelled if no longer needed.
+    3. "Acquired": The resources have been acquired by a caller to use for scheduling
+       remote Ray objects. Note that it is the responsibility of the caller to
+       schedule the Ray objects with these resources.
+       The associated resource request has been completed and can no longer be
+       cancelled. The acquired resources can be freed by the resource manager when
+       they are no longer used.
+
+    The flow is as follows:
+
+    .. code-block:: python
+
+        # Create resource manager
+        resource_manager = ResourceManager()
+
+        # Create resource request
+        resource_request = ResourceRequest([{"CPU": 4}])
+
+        # Pass to resource manager
+        resource_manager.request_resources(resource_request)
+
+        # Wait until ready
+        while not resource_manager.has_resources_ready(resource_request):
+            time.sleep(1)
+
+        # Once ready, acquire resources
+        acquired_resource = resource_manager.acquire_resources(resource_request)
+
+        # Bind to remote task or actor
+        annotated_remote_fn = acquired_resource.annotate_remote_entities(
+            [remote_fn])
+
+        # Run remote function. This will use the acquired resources
+        ray.get(annotated_remote_fn.remote())
+
+        # After using the resources, free
+        resource_manager.free_resources(annotated_resources)
+
+    """
+
+    def request_resources(self, resource_request: ResourceRequest):
+        """Request resources.
+
+        Depending on the backend, resources can trigger autoscaling. Requested
+        resources can be ready or not ready. Once they are "ready", they can
+        be acquired and used by remote Ray objects.
+
+        Resource requests can be cancelled anytime using ``cancel_resource_request()``.
+        Once acquired, the resource request is removed. Acquired resources can be
+        freed with ``free_resources()``.
+        """
+        raise NotImplementedError
+
+    def cancel_resource_request(self, resource_request: ResourceRequest):
+        """Cancel resource request.
+
+        Resource requests can be cancelled anytime before a resource is acquired.
+        Acquiring a resource will remove the associated resource request.
+        Acquired resources can be freed with ``free_resources()``.
+        """
+        raise NotImplementedError
+
+    def has_resources_ready(self, resource_request: ResourceRequest) -> bool:
+        """Returns True if resources for the given request are ready to be acquired."""
+        raise NotImplementedError
+
+    def acquire_resources(
+        self, resource_request: ResourceRequest
+    ) -> Optional[AcquiredResources]:
+        """Acquire resources. Returns None if resources are not ready to be acquired.
+
+        Acquiring resources will remove the associated resource request.
+        Acquired resources can be returned with ``free_resources()``.
+        """
+        raise NotImplementedError
+
+    def free_resources(self, acquired_resource: AcquiredResources):
+        """Free acquired resources from usage and return them to the resource manager.
+
+        Freeing resources will return the resources to the manager, but there are
+        no guarantees about the tasks and actors scheduled on the resources. The caller
+        should make sure that any references to tasks or actors scheduled on the
+        resources have been removed before calling ``free_resources()``.
+        """
+        raise NotImplementedError
+
+    def get_resource_futures(self) -> List[ray.ObjectRef]:
+        """Return futures for resources to await.
+
+        Depending on the backend, we use resource futures to determine availability
+        of resources (e.g. placement groups) or resolution of requests.
+        In this case, the futures can be awaited externally by the caller.
+
+        When a resource future resolved, the caller may call ``update_state()``
+        to force the resource manager to update its internal state immediately.
+        """
+        return []
+
+    def update_state(self):
+        """Update internal state of the resource manager.
+
+        The resource manager may have internal state that needs periodic updating.
+        For instance, depending on the backend, resource futures can be awaited
+        externally (with ``get_resource_futures()``).
+
+        If such a future resolved, the caller can instruct the resource
+        manager to update its internal state immediately.
+        """
+        pass
+
+    def clear(self):
+        """Reset internal state and clear all resources.
+
+        Calling this method will reset the resource manager to its initialization state.
+        All resources will be removed.
+
+        Clearing the state will remove tracked resources from the manager, but there are
+        no guarantees about the tasks and actors scheduled on the resources. The caller
+        should make sure that any references to tasks or actors scheduled on the
+        resources have been removed before calling ``clear()``.
+        """
+        raise NotImplementedError
+
+    def __reduce__(self):
+        """We disallow serialization.
+
+        Shared resource managers should live on an actor.
+        """
+        raise ValueError(
+            f"Resource managers cannot be serialized. Resource manager: {str(self)}"
+        )

python/ray/air/integrations/comet.py

@@ -0,0 +1,260 @@
+import os
+from pathlib import Path
+from typing import Dict, List
+
+import pyarrow.fs
+
+from ray.tune.experiment import Trial
+from ray.tune.logger import LoggerCallback
+from ray.tune.utils import flatten_dict
+
+
+def _import_comet():
+    """Try importing comet_ml.
+
+    Used to check if comet_ml is installed and, otherwise, pass an informative
+    error message.
+    """
+    if "COMET_DISABLE_AUTO_LOGGING" not in os.environ:
+        os.environ["COMET_DISABLE_AUTO_LOGGING"] = "1"
+
+    try:
+        import comet_ml  # noqa: F401
+    except ImportError:
+        raise RuntimeError("pip install 'comet-ml' to use CometLoggerCallback")
+
+    return comet_ml
+
+
+class CometLoggerCallback(LoggerCallback):
+    """CometLoggerCallback for logging Tune results to Comet.
+
+    Comet (https://comet.ml/site/) is a tool to manage and optimize the
+    entire ML lifecycle, from experiment tracking, model optimization
+    and dataset versioning to model production monitoring.
+
+    This Ray Tune ``LoggerCallback`` sends metrics and parameters to
+    Comet for tracking.
+
+    In order to use the CometLoggerCallback you must first install Comet
+    via ``pip install comet_ml``
+
+    Then set the following environment variables
+    ``export COMET_API_KEY=<Your API Key>``
+
+    Alternatively, you can also pass in your API Key as an argument to the
+    CometLoggerCallback constructor.
+
+    ``CometLoggerCallback(api_key=<Your API Key>)``
+
+    Args:
+            online: Whether to make use of an Online or
+                Offline Experiment. Defaults to True.
+            tags: Tags to add to the logged Experiment.
+                Defaults to None.
+            save_checkpoints: If ``True``, model checkpoints will be saved to
+                Comet ML as artifacts. Defaults to ``False``.
+            **experiment_kwargs: Other keyword arguments will be passed to the
+                constructor for comet_ml.Experiment (or OfflineExperiment if
+                online=False).
+
+    Please consult the Comet ML documentation for more information on the
+    Experiment and OfflineExperiment classes: https://comet.ml/site/
+
+    Example:
+
+    .. code-block:: python
+
+        from ray.air.integrations.comet import CometLoggerCallback
+        tune.run(
+            train,
+            config=config
+            callbacks=[CometLoggerCallback(
+                True,
+                ['tag1', 'tag2'],
+                workspace='my_workspace',
+                project_name='my_project_name'
+                )]
+        )
+
+    """
+
+    # Do not enable these auto log options unless overridden
+    _exclude_autolog = [
+        "auto_output_logging",
+        "log_git_metadata",
+        "log_git_patch",
+        "log_env_cpu",
+        "log_env_gpu",
+    ]
+
+    # Do not log these metrics.
+    _exclude_results = ["done", "should_checkpoint"]
+
+    # These values should be logged as system info instead of metrics.
+    _system_results = ["node_ip", "hostname", "pid", "date"]
+
+    # These values should be logged as "Other" instead of as metrics.
+    _other_results = ["trial_id", "experiment_id", "experiment_tag"]
+
+    _episode_results = ["hist_stats/episode_reward", "hist_stats/episode_lengths"]
+
+    def __init__(
+        self,
+        online: bool = True,
+        tags: List[str] = None,
+        save_checkpoints: bool = False,
+        **experiment_kwargs,
+    ):
+        _import_comet()
+        self.online = online
+        self.tags = tags
+        self.save_checkpoints = save_checkpoints
+        self.experiment_kwargs = experiment_kwargs
+
+        # Disable the specific autologging features that cause throttling.
+        self._configure_experiment_defaults()
+
+        # Mapping from trial to experiment object.
+        self._trial_experiments = {}
+
+        self._to_exclude = self._exclude_results.copy()
+        self._to_system = self._system_results.copy()
+        self._to_other = self._other_results.copy()
+        self._to_episodes = self._episode_results.copy()
+
+    def _configure_experiment_defaults(self):
+        """Disable the specific autologging features that cause throttling."""
+        for option in self._exclude_autolog:
+            if not self.experiment_kwargs.get(option):
+                self.experiment_kwargs[option] = False
+
+    def _check_key_name(self, key: str, item: str) -> bool:
+        """
+        Check if key argument is equal to item argument or starts with item and
+        a forward slash. Used for parsing trial result dictionary into ignored
+        keys, system metrics, episode logs, etc.
+        """
+        return key.startswith(item + "/") or key == item
+
+    def log_trial_start(self, trial: "Trial"):
+        """
+        Initialize an Experiment (or OfflineExperiment if self.online=False)
+        and start logging to Comet.
+
+        Args:
+            trial: Trial object.
+
+        """
+        _import_comet()  # is this necessary?
+        from comet_ml import Experiment, OfflineExperiment
+        from comet_ml.config import set_global_experiment
+
+        if trial not in self._trial_experiments:
+            experiment_cls = Experiment if self.online else OfflineExperiment
+            experiment = experiment_cls(**self.experiment_kwargs)
+            self._trial_experiments[trial] = experiment
+            # Set global experiment to None to allow for multiple experiments.
+            set_global_experiment(None)
+        else:
+            experiment = self._trial_experiments[trial]
+
+        experiment.set_name(str(trial))
+        experiment.add_tags(self.tags)
+        experiment.log_other("Created from", "Ray")
+
+        config = trial.config.copy()
+        config.pop("callbacks", None)
+        experiment.log_parameters(config)
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        """
+        Log the current result of a Trial upon each iteration.
+        """
+        if trial not in self._trial_experiments:
+            self.log_trial_start(trial)
+        experiment = self._trial_experiments[trial]
+        step = result["training_iteration"]
+
+        config_update = result.pop("config", {}).copy()
+        config_update.pop("callbacks", None)  # Remove callbacks
+        for k, v in config_update.items():
+            if isinstance(v, dict):
+                experiment.log_parameters(flatten_dict({k: v}, "/"), step=step)
+
+            else:
+                experiment.log_parameter(k, v, step=step)
+
+        other_logs = {}
+        metric_logs = {}
+        system_logs = {}
+        episode_logs = {}
+
+        flat_result = flatten_dict(result, delimiter="/")
+        for k, v in flat_result.items():
+            if any(self._check_key_name(k, item) for item in self._to_exclude):
+                continue
+
+            if any(self._check_key_name(k, item) for item in self._to_other):
+                other_logs[k] = v
+
+            elif any(self._check_key_name(k, item) for item in self._to_system):
+                system_logs[k] = v
+
+            elif any(self._check_key_name(k, item) for item in self._to_episodes):
+                episode_logs[k] = v
+
+            else:
+                metric_logs[k] = v
+
+        experiment.log_others(other_logs)
+        experiment.log_metrics(metric_logs, step=step)
+
+        for k, v in system_logs.items():
+            experiment.log_system_info(k, v)
+
+        for k, v in episode_logs.items():
+            experiment.log_curve(k, x=range(len(v)), y=v, step=step)
+
+    def log_trial_save(self, trial: "Trial"):
+        comet_ml = _import_comet()
+
+        if self.save_checkpoints and trial.checkpoint:
+            experiment = self._trial_experiments[trial]
+
+            artifact = comet_ml.Artifact(
+                name=f"checkpoint_{(str(trial))}", artifact_type="model"
+            )
+
+            checkpoint_root = None
+
+            if isinstance(trial.checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
+                checkpoint_root = trial.checkpoint.path
+                # Todo: For other filesystems, we may want to use
+                # artifact.add_remote() instead. However, this requires a full
+                # URI. We can add this once we have a way to retrieve it.
+
+            # Walk through checkpoint directory and add all files to artifact
+            if checkpoint_root:
+                for root, dirs, files in os.walk(checkpoint_root):
+                    rel_root = os.path.relpath(root, checkpoint_root)
+                    for file in files:
+                        local_file = Path(checkpoint_root, rel_root, file).as_posix()
+                        logical_path = Path(rel_root, file).as_posix()
+
+                        # Strip leading `./`
+                        if logical_path.startswith("./"):
+                            logical_path = logical_path[2:]
+
+                        artifact.add(local_file, logical_path=logical_path)
+
+            experiment.log_artifact(artifact)
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        self._trial_experiments[trial].end()
+        del self._trial_experiments[trial]
+
+    def __del__(self):
+        for trial, experiment in self._trial_experiments.items():
+            experiment.end()
+        self._trial_experiments = {}

python/ray/air/integrations/keras.py

@@ -0,0 +1,185 @@
+import shutil
+from typing import Dict, List, Optional, Union
+
+from tensorflow.keras.callbacks import Callback as KerasCallback
+
+import ray
+from ray.train.tensorflow import TensorflowCheckpoint
+from ray.util.annotations import PublicAPI
+
+
+class _Callback(KerasCallback):
+    """Base class for Air's Keras callbacks."""
+
+    _allowed = [
+        "epoch_begin",
+        "epoch_end",
+        "train_batch_begin",
+        "train_batch_end",
+        "test_batch_begin",
+        "test_batch_end",
+        "predict_batch_begin",
+        "predict_batch_end",
+        "train_begin",
+        "train_end",
+        "test_begin",
+        "test_end",
+        "predict_begin",
+        "predict_end",
+    ]
+
+    def __init__(self, on: Union[str, List[str]] = "validation_end"):
+        super(_Callback, self).__init__()
+
+        if not isinstance(on, list):
+            on = [on]
+        if any(w not in self._allowed for w in on):
+            raise ValueError(
+                "Invalid trigger time selected: {}. Must be one of {}".format(
+                    on, self._allowed
+                )
+            )
+        self._on = on
+
+    def _handle(self, logs: Dict, when: str):
+        raise NotImplementedError
+
+    def on_epoch_begin(self, epoch, logs=None):
+        if "epoch_begin" in self._on:
+            self._handle(logs, "epoch_begin")
+
+    def on_epoch_end(self, epoch, logs=None):
+        if "epoch_end" in self._on:
+            self._handle(logs, "epoch_end")
+
+    def on_train_batch_begin(self, batch, logs=None):
+        if "train_batch_begin" in self._on:
+            self._handle(logs, "train_batch_begin")
+
+    def on_train_batch_end(self, batch, logs=None):
+        if "train_batch_end" in self._on:
+            self._handle(logs, "train_batch_end")
+
+    def on_test_batch_begin(self, batch, logs=None):
+        if "test_batch_begin" in self._on:
+            self._handle(logs, "test_batch_begin")
+
+    def on_test_batch_end(self, batch, logs=None):
+        if "test_batch_end" in self._on:
+            self._handle(logs, "test_batch_end")
+
+    def on_predict_batch_begin(self, batch, logs=None):
+        if "predict_batch_begin" in self._on:
+            self._handle(logs, "predict_batch_begin")
+
+    def on_predict_batch_end(self, batch, logs=None):
+        if "predict_batch_end" in self._on:
+            self._handle(logs, "predict_batch_end")
+
+    def on_train_begin(self, logs=None):
+        if "train_begin" in self._on:
+            self._handle(logs, "train_begin")
+
+    def on_train_end(self, logs=None):
+        if "train_end" in self._on:
+            self._handle(logs, "train_end")
+
+    def on_test_begin(self, logs=None):
+        if "test_begin" in self._on:
+            self._handle(logs, "test_begin")
+
+    def on_test_end(self, logs=None):
+        if "test_end" in self._on:
+            self._handle(logs, "test_end")
+
+    def on_predict_begin(self, logs=None):
+        if "predict_begin" in self._on:
+            self._handle(logs, "predict_begin")
+
+    def on_predict_end(self, logs=None):
+        if "predict_end" in self._on:
+            self._handle(logs, "predict_end")
+
+
+@PublicAPI(stability="alpha")
+class ReportCheckpointCallback(_Callback):
+    """Keras callback for Ray Train reporting and checkpointing.
+
+    .. note::
+        Metrics are always reported with checkpoints, even if the event isn't specified
+        in ``report_metrics_on``.
+
+    Example:
+        .. code-block:: python
+
+            ############# Using it in TrainSession ###############
+            from ray.air.integrations.keras import ReportCheckpointCallback
+            def train_loop_per_worker():
+                strategy = tf.distribute.MultiWorkerMirroredStrategy()
+                with strategy.scope():
+                    model = build_model()
+
+                model.fit(dataset_shard, callbacks=[ReportCheckpointCallback()])
+
+    Args:
+        checkpoint_on: When to save checkpoints. Must be one of the Keras event hooks
+            (less the ``on_``), e.g. "train_start" or "predict_end". Defaults to
+            "epoch_end".
+        report_metrics_on: When to report metrics. Must be one of
+            the Keras event hooks (less the ``on_``), e.g.
+            "train_start" or "predict_end". Defaults to "epoch_end".
+        metrics: Metrics to report. If this is a list, each item describes
+            the metric key reported to Keras, and it's reported under the
+            same name. If this is a dict, each key is the name reported
+            and the respective value is the metric key reported to Keras.
+            If this is None, all Keras logs are reported.
+    """
+
+    def __init__(
+        self,
+        checkpoint_on: Union[str, List[str]] = "epoch_end",
+        report_metrics_on: Union[str, List[str]] = "epoch_end",
+        metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
+    ):
+        if isinstance(checkpoint_on, str):
+            checkpoint_on = [checkpoint_on]
+        if isinstance(report_metrics_on, str):
+            report_metrics_on = [report_metrics_on]
+
+        on = list(set(checkpoint_on + report_metrics_on))
+        super().__init__(on=on)
+
+        self._checkpoint_on: List[str] = checkpoint_on
+        self._report_metrics_on: List[str] = report_metrics_on
+        self._metrics = metrics
+
+    def _handle(self, logs: Dict, when: str):
+        assert when in self._checkpoint_on or when in self._report_metrics_on
+
+        metrics = self._get_reported_metrics(logs)
+
+        should_checkpoint = when in self._checkpoint_on
+        if should_checkpoint:
+            checkpoint = TensorflowCheckpoint.from_model(self.model)
+            ray.train.report(metrics, checkpoint=checkpoint)
+            # Clean up temporary checkpoint
+            shutil.rmtree(checkpoint.path, ignore_errors=True)
+        else:
+            ray.train.report(metrics, checkpoint=None)
+
+    def _get_reported_metrics(self, logs: Dict) -> Dict:
+        assert isinstance(self._metrics, (type(None), str, list, dict))
+
+        if self._metrics is None:
+            reported_metrics = logs
+        elif isinstance(self._metrics, str):
+            reported_metrics = {self._metrics: logs[self._metrics]}
+        elif isinstance(self._metrics, list):
+            reported_metrics = {metric: logs[metric] for metric in self._metrics}
+        elif isinstance(self._metrics, dict):
+            reported_metrics = {
+                key: logs[metric] for key, metric in self._metrics.items()
+            }
+
+        assert isinstance(reported_metrics, dict)
+        return reported_metrics

python/ray/air/integrations/mlflow.py

@@ -0,0 +1,343 @@
+import logging
+from types import ModuleType
+from typing import Dict, Optional, Union
+
+import ray
+from ray.air._internal import usage as air_usage
+from ray.air._internal.mlflow import _MLflowLoggerUtil
+from ray.air.constants import TRAINING_ITERATION
+from ray.tune.experiment import Trial
+from ray.tune.logger import LoggerCallback
+from ray.tune.result import TIMESTEPS_TOTAL
+from ray.tune.trainable.trainable_fn_utils import _in_tune_session
+from ray.util.annotations import PublicAPI
+
+try:
+    import mlflow
+except ImportError:
+    mlflow = None
+
+
+logger = logging.getLogger(__name__)
+
+
+class _NoopModule:
+    def __getattr__(self, item):
+        return _NoopModule()
+
+    def __call__(self, *args, **kwargs):
+        return None
+
+
+@PublicAPI(stability="alpha")
+def setup_mlflow(
+    config: Optional[Dict] = None,
+    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,
+    run_name: Optional[str] = None,
+    create_experiment_if_not_exists: bool = False,
+    tags: Optional[Dict] = None,
+    rank_zero_only: bool = True,
+) -> Union[ModuleType, _NoopModule]:
+    """Set up a MLflow session.
+
+    This function can be used to initialize an MLflow session in a
+    (distributed) training or tuning run. The session will be created on the trainable.
+
+    By default, the MLflow experiment ID is the Ray trial ID and the
+    MLlflow experiment name is the Ray trial name. These settings can be overwritten by
+    passing the respective keyword arguments.
+
+    The ``config`` dict is automatically logged as the run parameters (excluding the
+    mlflow settings).
+
+    In distributed training with Ray Train, only the zero-rank worker will initialize
+    mlflow. All other workers will return a noop client, so that logging is not
+    duplicated in a distributed run. This can be disabled by passing
+    ``rank_zero_only=False``, which will then initialize mlflow in every training
+    worker. Note: for Ray Tune, there's no concept of worker ranks, so the `rank_zero_only` is ignored.
+
+    This function will return the ``mlflow`` module or a noop module for
+    non-rank zero workers ``if rank_zero_only=True``. By using
+    ``mlflow = setup_mlflow(config)`` you can ensure that only the rank zero worker
+    calls the mlflow API.
+
+    Args:
+        config: Configuration dict to be logged to mlflow as parameters.
+        tracking_uri: The tracking URI for MLflow tracking. If using
+            Tune in a multi-node setting, make sure to use a remote server for
+            tracking.
+        registry_uri: The registry URI for the MLflow model registry.
+        experiment_id: The id of an already created MLflow experiment.
+            All logs from all trials in ``tune.Tuner()`` will be reported to this
+            experiment. If this is not provided or the experiment with this
+            id does not exist, you must provide an``experiment_name``. This
+            parameter takes precedence over ``experiment_name``.
+        experiment_name: The name of an already existing MLflow
+            experiment. All logs from all trials in ``tune.Tuner()`` will be
+            reported to this experiment. If this is not provided, you must
+            provide a valid ``experiment_id``.
+        tracking_token: A token to use for HTTP authentication when
+            logging to a remote tracking server. This is useful when you
+            want to log to a Databricks server, for example. This value will
+            be used to set the MLFLOW_TRACKING_TOKEN environment variable on
+            all the remote training processes.
+        artifact_location: The location to store run artifacts.
+            If not provided, MLFlow picks an appropriate default.
+            Ignored if experiment already exists.
+        run_name: Name of the new MLflow run that will be created.
+            If not set, will default to the ``experiment_name``.
+        create_experiment_if_not_exists: Whether to create an
+            experiment with the provided name if it does not already
+            exist. Defaults to False.
+        tags: Tags to set for the new run.
+        rank_zero_only: If True, will return an initialized session only for the
+            rank 0 worker in distributed training. If False, will initialize a
+            session for all workers. Defaults to True.
+
+    Example:
+
+        Per default, you can just call ``setup_mlflow`` and continue to use
+        MLflow like you would normally do:
+
+        .. code-block:: python
+
+            from ray.air.integrations.mlflow import setup_mlflow
+
+            def training_loop(config):
+                mlflow = setup_mlflow(config)
+                # ...
+                mlflow.log_metric(key="loss", val=0.123, step=0)
+
+        In distributed data parallel training, you can utilize the return value of
+        ``setup_mlflow``. This will make sure it is only invoked on the first worker
+        in distributed training runs.
+
+        .. code-block:: python
+
+            from ray.air.integrations.mlflow import setup_mlflow
+
+            def training_loop(config):
+                mlflow = setup_mlflow(config)
+                # ...
+                mlflow.log_metric(key="loss", val=0.123, step=0)
+
+
+        You can also use MlFlow's autologging feature if using a training
+        framework like Pytorch Lightning, XGBoost, etc. More information can be
+        found here
+        (https://mlflow.org/docs/latest/tracking.html#automatic-logging).
+
+        .. code-block:: python
+
+            from ray.air.integrations.mlflow import setup_mlflow
+
+            def train_fn(config):
+                mlflow = setup_mlflow(config)
+                mlflow.autolog()
+                xgboost_results = xgb.train(config, ...)
+
+    Returns:
+        The ``mlflow`` module, or a noop module for non-rank-zero workers when
+        ``rank_zero_only`` is True.
+    """
+    if not mlflow:
+        raise RuntimeError(
+            "mlflow was not found - please install with `pip install mlflow`"
+        )
+
+    default_trial_id = None
+    default_trial_name = None
+
+    try:
+        if _in_tune_session():
+            context: ray.tune.TuneContext = ray.tune.get_context()
+            default_trial_id = context.get_trial_id()
+            default_trial_name = context.get_trial_name()
+        else:
+            context: ray.train.TrainContext = ray.train.get_context()
+            if rank_zero_only and context.get_world_rank() != 0:
+                return _NoopModule()
+    except RuntimeError:
+        default_trial_id = None
+        default_trial_name = None
+
+    _config = config.copy() if config else {}
+
+    experiment_id = experiment_id or default_trial_id
+    experiment_name = experiment_name or default_trial_name
+
+    # Setup mlflow
+    mlflow_util = _MLflowLoggerUtil()
+    mlflow_util.setup_mlflow(
+        tracking_uri=tracking_uri,
+        registry_uri=registry_uri,
+        experiment_id=experiment_id,
+        experiment_name=experiment_name,
+        tracking_token=tracking_token,
+        artifact_location=artifact_location,
+        create_experiment_if_not_exists=create_experiment_if_not_exists,
+    )
+
+    mlflow_util.start_run(
+        run_name=run_name or experiment_name,
+        tags=tags,
+        set_active=True,
+    )
+    mlflow_util.log_params(_config)
+
+    # Record `setup_mlflow` usage when everything has setup successfully.
+    air_usage.tag_setup_mlflow()
+
+    return mlflow_util._mlflow
+
+
+class MLflowLoggerCallback(LoggerCallback):
+    """MLflow Logger to automatically log Tune results and config to MLflow.
+
+    MLflow (https://mlflow.org) Tracking is an open source library for
+    recording and querying experiments. This Ray Tune ``LoggerCallback``
+    sends information (config parameters, training results & metrics,
+    and artifacts) to MLflow for automatic experiment tracking.
+
+    Keep in mind that the callback will open an MLflow session on the driver and
+    not on the trainable. Therefore, it is not possible to call MLflow functions
+    like ``mlflow.log_figure()`` inside the trainable as there is no MLflow session
+    on the trainable. For more fine grained control, use
+    :func:`ray.air.integrations.mlflow.setup_mlflow`.
+
+    Args:
+        tracking_uri: The tracking URI for where to manage experiments
+            and runs. This can either be a local file path or a remote server.
+            This arg gets passed directly to mlflow
+            initialization. When using Tune in a multi-node setting, make sure
+            to set this to a remote server and not a local file path.
+        registry_uri: The registry URI that gets passed directly to
+            mlflow initialization.
+        experiment_name: The experiment name to use for this Tune run.
+            If the experiment with the name already exists with MLflow,
+            it will be reused. If not, a new experiment will be created with
+            that name.
+        tags: An optional dictionary of string keys and values to set
+            as tags on the run
+        tracking_token: Tracking token used to authenticate with MLflow.
+        save_artifact: If set to True, automatically save the entire
+            contents of the Tune local_dir as an artifact to the
+            corresponding run in MlFlow.
+        log_params_on_trial_end: If set to True, log parameters to MLflow
+            at the end of the trial instead of at the beginning
+
+    Example:
+
+    .. code-block:: python
+
+        from ray.air.integrations.mlflow import MLflowLoggerCallback
+
+        tags = { "user_name" : "John",
+                 "git_commit_hash" : "abc123"}
+
+        tune.run(
+            train_fn,
+            config={
+                # define search space here
+                "parameter_1": tune.choice([1, 2, 3]),
+                "parameter_2": tune.choice([4, 5, 6]),
+            },
+            callbacks=[MLflowLoggerCallback(
+                experiment_name="experiment1",
+                tags=tags,
+                save_artifact=True,
+                log_params_on_trial_end=True)])
+
+    """
+
+    def __init__(
+        self,
+        tracking_uri: Optional[str] = None,
+        *,
+        registry_uri: Optional[str] = None,
+        experiment_name: Optional[str] = None,
+        tags: Optional[Dict] = None,
+        tracking_token: Optional[str] = None,
+        save_artifact: bool = False,
+        log_params_on_trial_end: bool = False,
+    ):
+
+        self.tracking_uri = tracking_uri
+        self.registry_uri = registry_uri
+        self.experiment_name = experiment_name
+        self.tags = tags
+        self.tracking_token = tracking_token
+        self.should_save_artifact = save_artifact
+        self.log_params_on_trial_end = log_params_on_trial_end
+
+        self.mlflow_util = _MLflowLoggerUtil()
+
+        if ray.util.client.ray.is_connected():
+            logger.warning(
+                "When using MLflowLoggerCallback with Ray Client, "
+                "it is recommended to use a remote tracking "
+                "server. If you are using a MLflow tracking server "
+                "backed by the local filesystem, then it must be "
+                "setup on the server side and not on the client "
+                "side."
+            )
+
+    def setup(self, *args, **kwargs):
+        # Setup the mlflow logging util.
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+            experiment_name=self.experiment_name,
+            tracking_token=self.tracking_token,
+        )
+
+        if self.tags is None:
+            # Create empty dictionary for tags if not given explicitly
+            self.tags = {}
+
+        self._trial_runs = {}
+
+    def log_trial_start(self, trial: "Trial"):
+        # Create run if not already exists.
+        if trial not in self._trial_runs:
+
+            # Set trial name in tags
+            tags = self.tags.copy()
+            tags["trial_name"] = str(trial)
+
+            run = self.mlflow_util.start_run(tags=tags, run_name=str(trial))
+            self._trial_runs[trial] = run.info.run_id
+
+        run_id = self._trial_runs[trial]
+
+        # Log the config parameters.
+        config = trial.config
+        if not self.log_params_on_trial_end:
+            self.mlflow_util.log_params(run_id=run_id, params_to_log=config)
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        step = result.get(TIMESTEPS_TOTAL) or result[TRAINING_ITERATION]
+        run_id = self._trial_runs[trial]
+        self.mlflow_util.log_metrics(run_id=run_id, metrics_to_log=result, step=step)
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        run_id = self._trial_runs[trial]
+
+        # Log the artifact if set_artifact is set to True.
+        if self.should_save_artifact:
+            self.mlflow_util.save_artifacts(run_id=run_id, dir=trial.local_path)
+
+        # Stop the run once trial finishes.
+        status = "FINISHED" if not failed else "FAILED"
+
+        # Log the config parameters.
+        config = trial.config
+        if self.log_params_on_trial_end:
+            self.mlflow_util.log_params(run_id=run_id, params_to_log=config)
+
+        self.mlflow_util.end_run(run_id=run_id, status=status)

python/ray/air/integrations/wandb.py

@@ -0,0 +1,810 @@
+import enum
+import os
+import pickle
+import urllib
+import warnings
+from numbers import Number
+from types import ModuleType
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import pyarrow.fs
+
+import ray
+from ray import logger
+from ray._common.utils import load_class
+from ray.air._internal import usage as air_usage
+from ray.air.constants import TRAINING_ITERATION
+from ray.air.util.node import _force_on_current_node
+from ray.train._internal.session import get_session
+from ray.train._internal.syncer import DEFAULT_SYNC_TIMEOUT
+from ray.tune.experiment import Trial
+from ray.tune.logger import LoggerCallback
+from ray.tune.utils import flatten_dict
+from ray.util import PublicAPI
+from ray.util.queue import Queue
+
+try:
+    import wandb
+    from wandb.sdk.data_types.base_types.wb_value import WBValue
+    from wandb.sdk.data_types.image import Image
+    from wandb.sdk.data_types.video import Video
+    from wandb.sdk.lib.disabled import RunDisabled
+    from wandb.util import json_dumps_safer
+    from wandb.wandb_run import Run
+except ImportError:
+    wandb = json_dumps_safer = Run = RunDisabled = WBValue = None
+
+
+WANDB_ENV_VAR = "WANDB_API_KEY"
+WANDB_PROJECT_ENV_VAR = "WANDB_PROJECT_NAME"
+WANDB_GROUP_ENV_VAR = "WANDB_GROUP_NAME"
+WANDB_MODE_ENV_VAR = "WANDB_MODE"
+# Hook that is invoked before wandb.init in the setup method of WandbLoggerCallback
+# to populate the API key if it isn't already set when initializing the callback.
+# It doesn't take in any arguments and returns the W&B API key.
+# Example: "your.module.wandb_setup_api_key_hook".
+WANDB_SETUP_API_KEY_HOOK = "WANDB_SETUP_API_KEY_HOOK"
+# Hook that is invoked before wandb.init in the setup method of WandbLoggerCallback
+# to populate environment variables to specify the location
+# (project and group) of the W&B run.
+# It doesn't take in any arguments and doesn't return anything, but it does populate
+# WANDB_PROJECT_NAME and WANDB_GROUP_NAME.
+# Example: "your.module.wandb_populate_run_location_hook".
+WANDB_POPULATE_RUN_LOCATION_HOOK = "WANDB_POPULATE_RUN_LOCATION_HOOK"
+# Hook that is invoked after running wandb.init in WandbLoggerCallback
+# to process information about the W&B run.
+# It takes in a W&B run object and doesn't return anything.
+# Example: "your.module.wandb_process_run_info_hook".
+WANDB_PROCESS_RUN_INFO_HOOK = "WANDB_PROCESS_RUN_INFO_HOOK"
+
+
+@PublicAPI(stability="alpha")
+def setup_wandb(
+    config: Optional[Dict] = None,
+    api_key: Optional[str] = None,
+    api_key_file: Optional[str] = None,
+    rank_zero_only: bool = True,
+    **kwargs,
+) -> Union[Run, RunDisabled]:
+    """Set up a Weights & Biases session.
+
+    This function can be used to initialize a Weights & Biases session in a
+    (distributed) training or tuning run.
+
+    By default, the run ID is the trial ID, the run name is the trial name, and
+    the run group is the experiment name. These settings can be overwritten by
+    passing the respective arguments as ``kwargs``, which will be passed to
+    ``wandb.init()``.
+
+    In distributed training with Ray Train, only the zero-rank worker will initialize
+    wandb. All other workers will return a disabled run object, so that logging is not
+    duplicated in a distributed run. This can be disabled by passing
+    ``rank_zero_only=False``, which will then initialize wandb in every training
+    worker.
+
+    The ``config`` argument will be passed to Weights and Biases and will be logged
+    as the run configuration.
+
+    If no API key or key file are passed, wandb will try to authenticate
+    using locally stored credentials, created for instance by running ``wandb login``.
+
+    Keyword arguments passed to ``setup_wandb()`` will be passed to
+    ``wandb.init()`` and take precedence over any potential default settings.
+
+    Args:
+        config: Configuration dict to be logged to Weights and Biases. Can contain
+            arguments for ``wandb.init()`` as well as authentication information.
+        api_key: API key to use for authentication with Weights and Biases.
+        api_key_file: File pointing to API key for with Weights and Biases.
+        rank_zero_only: If True, will return an initialized session only for the
+            rank 0 worker in distributed training. If False, will initialize a
+            session for all workers.
+        **kwargs: Passed to ``wandb.init()``.
+
+    Example:
+
+        .. code-block:: python
+
+            from ray.air.integrations.wandb import setup_wandb
+
+            def training_loop(config):
+                wandb = setup_wandb(config)
+                # ...
+                wandb.log({"loss": 0.123})
+
+    Returns:
+        The initialized wandb run, or a disabled run for non-rank-zero workers
+        when ``rank_zero_only`` is True.
+    """
+    if not wandb:
+        raise RuntimeError(
+            "Wandb was not found - please install with `pip install wandb`"
+        )
+
+    default_trial_id = None
+    default_trial_name = None
+    default_experiment_name = None
+
+    # Do a try-catch here if we are not in a train session
+    session = get_session()
+
+    if rank_zero_only:
+        # Check if we are in a train session and if we are not the rank 0 worker
+        if session and session.world_rank is not None and session.world_rank != 0:
+            return RunDisabled()
+
+    if session:
+        default_trial_id = session.trial_id
+        default_trial_name = session.trial_name
+        default_experiment_name = session.experiment_name
+
+    # Default init kwargs
+    wandb_init_kwargs = {
+        "trial_id": kwargs.get("trial_id") or default_trial_id,
+        "trial_name": kwargs.get("trial_name") or default_trial_name,
+        "group": kwargs.get("group") or default_experiment_name,
+    }
+    # Passed kwargs take precedence over default kwargs
+    wandb_init_kwargs.update(kwargs)
+
+    return _setup_wandb(
+        config=config, api_key=api_key, api_key_file=api_key_file, **wandb_init_kwargs
+    )
+
+
+def _setup_wandb(
+    trial_id: str,
+    trial_name: str,
+    config: Optional[Dict] = None,
+    api_key: Optional[str] = None,
+    api_key_file: Optional[str] = None,
+    _wandb: Optional[ModuleType] = None,
+    **kwargs,
+) -> Union[Run, RunDisabled]:
+    _config = config.copy() if config else {}
+
+    # If key file is specified, set
+    if api_key_file:
+        api_key_file = os.path.expanduser(api_key_file)
+
+    _set_api_key(api_key_file, api_key)
+    project = _get_wandb_project(kwargs.pop("project", None))
+    group = kwargs.pop("group", os.environ.get(WANDB_GROUP_ENV_VAR))
+
+    # Remove unpickleable items.
+    _config = _clean_log(_config)
+
+    wandb_init_kwargs = dict(
+        id=trial_id,
+        name=trial_name,
+        resume=True,
+        reinit=True,
+        allow_val_change=True,
+        config=_config,
+        project=project,
+        group=group,
+    )
+
+    # Update config (e.g. set any other parameters in the call to wandb.init)
+    wandb_init_kwargs.update(**kwargs)
+
+    # On windows, we can't fork
+    if os.name == "nt":
+        os.environ["WANDB_START_METHOD"] = "thread"
+    else:
+        os.environ["WANDB_START_METHOD"] = "fork"
+
+    _wandb = _wandb or wandb
+
+    run = _wandb.init(**wandb_init_kwargs)
+    _run_wandb_process_run_info_hook(run)
+
+    # Record `setup_wandb` usage when everything has setup successfully.
+    air_usage.tag_setup_wandb()
+
+    return run
+
+
+def _is_allowed_type(obj):
+    """Return True if type is allowed for logging to wandb"""
+    if isinstance(obj, np.ndarray) and obj.size == 1:
+        return isinstance(obj.item(), Number)
+    if isinstance(obj, Sequence) and len(obj) > 0:
+        return isinstance(obj[0], (Image, Video, WBValue))
+    return isinstance(obj, (Number, WBValue))
+
+
+def _clean_log(
+    obj: Any,
+    *,
+    video_kwargs: Optional[Dict[str, Any]] = None,
+    image_kwargs: Optional[Dict[str, Any]] = None,
+):
+    # Fixes https://github.com/ray-project/ray/issues/10631
+    if video_kwargs is None:
+        video_kwargs = {}
+    if image_kwargs is None:
+        image_kwargs = {}
+    if isinstance(obj, dict):
+        return {
+            k: _clean_log(v, video_kwargs=video_kwargs, image_kwargs=image_kwargs)
+            for k, v in obj.items()
+        }
+    elif isinstance(obj, (list, set)):
+        return [
+            _clean_log(v, video_kwargs=video_kwargs, image_kwargs=image_kwargs)
+            for v in obj
+        ]
+    elif isinstance(obj, tuple):
+        return tuple(
+            _clean_log(v, video_kwargs=video_kwargs, image_kwargs=image_kwargs)
+            for v in obj
+        )
+    elif isinstance(obj, np.ndarray) and obj.ndim == 3:
+        # Must be single image (H, W, C).
+        return Image(obj, **image_kwargs)
+    elif isinstance(obj, np.ndarray) and obj.ndim == 4:
+        # Must be batch of images (N >= 1, H, W, C).
+        return (
+            _clean_log(
+                [Image(v, **image_kwargs) for v in obj],
+                video_kwargs=video_kwargs,
+                image_kwargs=image_kwargs,
+            )
+            if obj.shape[0] > 1
+            else Image(obj[0], **image_kwargs)
+        )
+    elif isinstance(obj, np.ndarray) and obj.ndim == 5:
+        # Must be batch of videos (N >= 1, T, C, W, H).
+        return (
+            _clean_log(
+                [Video(v, **video_kwargs) for v in obj],
+                video_kwargs=video_kwargs,
+                image_kwargs=image_kwargs,
+            )
+            if obj.shape[0] > 1
+            else Video(obj[0], **video_kwargs)
+        )
+    elif _is_allowed_type(obj):
+        return obj
+
+    # Else
+
+    try:
+        # This is what wandb uses internally. If we cannot dump
+        # an object using this method, wandb will raise an exception.
+        json_dumps_safer(obj)
+
+        # This is probably unnecessary, but left here to be extra sure.
+        pickle.dumps(obj)
+
+        return obj
+    except Exception:
+        # give up, similar to _SafeFallBackEncoder
+        fallback = str(obj)
+
+        # Try to convert to int
+        try:
+            fallback = int(fallback)
+            return fallback
+        except ValueError:
+            pass
+
+        # Try to convert to float
+        try:
+            fallback = float(fallback)
+            return fallback
+        except ValueError:
+            pass
+
+        # Else, return string
+        return fallback
+
+
+def _get_wandb_project(project: Optional[str] = None) -> Optional[str]:
+    """Get W&B project from environment variable or external hook if not passed
+    as and argument."""
+    if (
+        not project
+        and not os.environ.get(WANDB_PROJECT_ENV_VAR)
+        and os.environ.get(WANDB_POPULATE_RUN_LOCATION_HOOK)
+    ):
+        # Try to populate WANDB_PROJECT_ENV_VAR and WANDB_GROUP_ENV_VAR
+        # from external hook
+        try:
+            load_class(os.environ[WANDB_POPULATE_RUN_LOCATION_HOOK])()
+        except Exception as e:
+            logger.exception(
+                f"Error executing {WANDB_POPULATE_RUN_LOCATION_HOOK} to "
+                f"populate {WANDB_PROJECT_ENV_VAR} and {WANDB_GROUP_ENV_VAR}: {e}",
+                exc_info=e,
+            )
+    if not project and os.environ.get(WANDB_PROJECT_ENV_VAR):
+        # Try to get project and group from environment variables if not
+        # passed through WandbLoggerCallback.
+        project = os.environ.get(WANDB_PROJECT_ENV_VAR)
+    return project
+
+
+def _set_api_key(api_key_file: Optional[str] = None, api_key: Optional[str] = None):
+    """Set WandB API key from `wandb_config`. Will pop the
+    `api_key_file` and `api_key` keys from `wandb_config` parameter.
+
+    The order of fetching the API key is:
+      1) From `api_key` or `api_key_file` arguments
+      2) From WANDB_API_KEY environment variables
+      3) User already logged in to W&B (wandb.api.api_key set)
+      4) From external hook WANDB_SETUP_API_KEY_HOOK
+    """
+    if os.environ.get(WANDB_MODE_ENV_VAR) in {"offline", "disabled"}:
+        return
+
+    if api_key_file:
+        if api_key:
+            raise ValueError("Both WandB `api_key_file` and `api_key` set.")
+        with open(api_key_file, "rt") as fp:
+            api_key = fp.readline().strip()
+
+    if not api_key and not os.environ.get(WANDB_ENV_VAR):
+        # Check if user is already logged into wandb.
+        try:
+            wandb.ensure_configured()
+            if wandb.api.api_key:
+                logger.info("Already logged into W&B.")
+                return
+        except AttributeError:
+            pass
+        # Try to get API key from external hook
+        if WANDB_SETUP_API_KEY_HOOK in os.environ:
+            try:
+                api_key = load_class(os.environ[WANDB_SETUP_API_KEY_HOOK])()
+            except Exception as e:
+                logger.exception(
+                    f"Error executing {WANDB_SETUP_API_KEY_HOOK} to setup API key: {e}",
+                    exc_info=e,
+                )
+    if api_key:
+        os.environ[WANDB_ENV_VAR] = api_key
+    elif not os.environ.get(WANDB_ENV_VAR):
+        raise ValueError(
+            "No WandB API key found. Either set the {} environment "
+            "variable, pass `api_key` or `api_key_file` to the"
+            "`WandbLoggerCallback` class as arguments, "
+            "or run `wandb login` from the command line".format(WANDB_ENV_VAR)
+        )
+
+
+def _run_wandb_process_run_info_hook(run: Any) -> None:
+    """Run external hook to process information about wandb run"""
+    if WANDB_PROCESS_RUN_INFO_HOOK in os.environ:
+        try:
+            load_class(os.environ[WANDB_PROCESS_RUN_INFO_HOOK])(run)
+        except Exception as e:
+            logger.exception(
+                f"Error calling {WANDB_PROCESS_RUN_INFO_HOOK}: {e}", exc_info=e
+            )
+
+
+class _QueueItem(enum.Enum):
+    END = enum.auto()
+    RESULT = enum.auto()
+    CHECKPOINT = enum.auto()
+
+
+class _WandbLoggingActor:
+    """
+    Wandb assumes that each trial's information should be logged from a
+    separate process. We use Ray actors as forking multiprocessing
+    processes is not supported by Ray and spawn processes run into pickling
+    problems.
+
+    We use a queue for the driver to communicate with the logging process.
+    The queue accepts the following items:
+
+    - If it's a dict, it is assumed to be a result and will be logged using
+      ``wandb.log()``
+    - If it's a checkpoint object, it will be saved using ``wandb.log_artifact()``.
+    """
+
+    def __init__(
+        self,
+        logdir: str,
+        queue: Queue,
+        exclude: List[str],
+        to_config: List[str],
+        *args,
+        **kwargs,
+    ):
+        import wandb
+
+        self._wandb = wandb
+
+        os.chdir(logdir)
+        self.queue = queue
+        self._exclude = set(exclude)
+        self._to_config = set(to_config)
+        self.args = args
+        self.kwargs = kwargs
+
+        self._trial_name = self.kwargs.get("name", "unknown")
+        self._logdir = logdir
+
+    def run(self):
+        # Since we're running in a separate process already, use threads.
+        os.environ["WANDB_START_METHOD"] = "thread"
+        run = self._wandb.init(*self.args, **self.kwargs)
+        run.config.trial_log_path = self._logdir
+
+        _run_wandb_process_run_info_hook(run)
+
+        while True:
+            item_type, item_content = self.queue.get()
+            if item_type == _QueueItem.END:
+                break
+
+            if item_type == _QueueItem.CHECKPOINT:
+                self._handle_checkpoint(item_content)
+                continue
+
+            assert item_type == _QueueItem.RESULT
+            log, config_update = self._handle_result(item_content)
+            try:
+                self._wandb.config.update(config_update, allow_val_change=True)
+                self._wandb.log(log, step=log.get(TRAINING_ITERATION))
+            except urllib.error.HTTPError as e:
+                # Ignore HTTPError. Missing a few data points is not a
+                # big issue, as long as things eventually recover.
+                logger.warning("Failed to log result to w&b: {}".format(str(e)))
+            except FileNotFoundError as e:
+                logger.error(
+                    "FileNotFoundError: Did not log result to Weights & Biases. "
+                    "Possible cause: relative file path used instead of absolute path. "
+                    "Error: %s",
+                    e,
+                )
+        self._wandb.finish()
+
+    def _handle_checkpoint(self, checkpoint_path: str):
+        artifact = self._wandb.Artifact(
+            name=f"checkpoint_{self._trial_name}", type="model"
+        )
+        artifact.add_dir(checkpoint_path)
+        self._wandb.log_artifact(artifact)
+
+    def _handle_result(self, result: Dict) -> Tuple[Dict, Dict]:
+        config_update = result.get("config", {}).copy()
+        log = {}
+        flat_result = flatten_dict(result, delimiter="/")
+
+        for k, v in flat_result.items():
+            if any(k.startswith(item + "/") or k == item for item in self._exclude):
+                continue
+            elif any(k.startswith(item + "/") or k == item for item in self._to_config):
+                config_update[k] = v
+            elif not _is_allowed_type(v):
+                continue
+            else:
+                log[k] = v
+
+        config_update.pop("callbacks", None)  # Remove callbacks
+        return log, config_update
+
+
+@PublicAPI(stability="alpha")
+class WandbLoggerCallback(LoggerCallback):
+    """WandbLoggerCallback
+
+    Weights and biases (https://www.wandb.ai/) is a tool for experiment
+    tracking, model optimization, and dataset versioning. This Ray Tune
+    ``LoggerCallback`` sends metrics to Wandb for automatic tracking and
+    visualization.
+
+    Example:
+
+        .. testcode::
+
+            import random
+
+            from ray import tune
+            from ray.air.integrations.wandb import WandbLoggerCallback
+
+
+            def train_func(config):
+                offset = random.random() / 5
+                for epoch in range(2, config["epochs"]):
+                    acc = 1 - (2 + config["lr"]) ** -epoch - random.random() / epoch - offset
+                    loss = (2 + config["lr"]) ** -epoch + random.random() / epoch + offset
+                    train.report({"acc": acc, "loss": loss})
+
+
+            tuner = tune.Tuner(
+                train_func,
+                param_space={
+                    "lr": tune.grid_search([0.001, 0.01, 0.1, 1.0]),
+                    "epochs": 10,
+                },
+                run_config=tune.RunConfig(
+                    callbacks=[WandbLoggerCallback(project="Optimization_Project")]
+                ),
+            )
+            results = tuner.fit()
+
+        .. testoutput::
+            :hide:
+
+            ...
+
+    Args:
+        project: Name of the Wandb project. Mandatory.
+        group: Name of the Wandb group. Defaults to the trainable
+            name.
+        api_key_file: Path to file containing the Wandb API KEY. This
+            file only needs to be present on the node running the Tune script
+            if using the WandbLogger.
+        api_key: Wandb API Key. Alternative to setting ``api_key_file``.
+        excludes: List of metrics and config that should be excluded from
+            the log.
+        log_config: Boolean indicating if the ``config`` parameter of
+            the ``results`` dict should be logged. This makes sense if
+            parameters will change during training, e.g. with
+            PopulationBasedTraining. Defaults to False.
+        upload_checkpoints: If ``True``, model checkpoints will be uploaded to
+            Wandb as artifacts. Defaults to ``False``.
+        save_checkpoints: Deprecated alias of ``upload_checkpoints``. Defaults to
+            ``False``.
+        upload_timeout: Maximum time in seconds to wait for pending uploads to
+            wandb when the experiment ends. Defaults to the Ray Train default
+            sync timeout.
+        video_kwargs: Dictionary of keyword arguments passed to wandb.Video()
+            when logging videos. Videos have to be logged as 5D numpy arrays
+            to be affected by this parameter. For valid keyword arguments, see
+            https://docs.wandb.ai/ref/python/data-types/video/. Defaults to ``None``.
+        image_kwargs: Dictionary of keyword arguments passed to wandb.Image()
+            when logging images. Images have to be logged as 3D or 4D numpy arrays
+            to be affected by this parameter. For valid keyword arguments, see
+            https://docs.wandb.ai/ref/python/data-types/image/. Defaults to ``None``.
+        **kwargs: The keyword arguments will be passed to ``wandb.init()``.
+
+    Wandb's ``group``, ``run_id`` and ``run_name`` are automatically selected
+    by Tune, but can be overwritten by filling out the respective configuration
+    values.
+
+    Please see here for all other valid configuration settings:
+    https://docs.wandb.ai/ref/python/init/
+    """  # noqa: E501
+
+    # Do not log these result keys
+    _exclude_results = ["done", "should_checkpoint"]
+
+    AUTO_CONFIG_KEYS = [
+        "trial_id",
+        "experiment_tag",
+        "node_ip",
+        "experiment_id",
+        "hostname",
+        "pid",
+        "date",
+    ]
+    """Results that are saved with `wandb.config` instead of `wandb.log`."""
+
+    _logger_actor_cls = _WandbLoggingActor
+
+    def __init__(
+        self,
+        project: Optional[str] = None,
+        group: Optional[str] = None,
+        api_key_file: Optional[str] = None,
+        api_key: Optional[str] = None,
+        excludes: Optional[List[str]] = None,
+        log_config: bool = False,
+        upload_checkpoints: bool = False,
+        save_checkpoints: bool = False,
+        upload_timeout: int = DEFAULT_SYNC_TIMEOUT,
+        video_kwargs: Optional[dict] = None,
+        image_kwargs: Optional[dict] = None,
+        **kwargs,
+    ):
+        if not wandb:
+            raise RuntimeError(
+                "Wandb was not found - please install with `pip install wandb`"
+            )
+
+        if save_checkpoints:
+            warnings.warn(
+                "`save_checkpoints` is deprecated. Use `upload_checkpoints` instead.",
+                DeprecationWarning,
+            )
+            upload_checkpoints = save_checkpoints
+
+        self.project = project
+        self.group = group
+        self.api_key_path = api_key_file
+        self.api_key = api_key
+        self.excludes = excludes or []
+        self.log_config = log_config
+        self.upload_checkpoints = upload_checkpoints
+        self._upload_timeout = upload_timeout
+        self.video_kwargs = video_kwargs or {}
+        self.image_kwargs = image_kwargs or {}
+        self.kwargs = kwargs
+
+        self._remote_logger_class = None
+
+        self._trial_logging_actors: Dict[
+            "Trial", ray.actor.ActorHandle[_WandbLoggingActor]
+        ] = {}
+        self._trial_logging_futures: Dict["Trial", ray.ObjectRef] = {}
+        self._logging_future_to_trial: Dict[ray.ObjectRef, "Trial"] = {}
+        self._trial_queues: Dict["Trial", Queue] = {}
+
+    def setup(self, *args, **kwargs):
+        self.api_key_file = (
+            os.path.expanduser(self.api_key_path) if self.api_key_path else None
+        )
+        _set_api_key(self.api_key_file, self.api_key)
+
+        self.project = _get_wandb_project(self.project)
+        if not self.project:
+            raise ValueError(
+                "Please pass the project name as argument or through "
+                f"the {WANDB_PROJECT_ENV_VAR} environment variable."
+            )
+        if not self.group and os.environ.get(WANDB_GROUP_ENV_VAR):
+            self.group = os.environ.get(WANDB_GROUP_ENV_VAR)
+
+    def log_trial_start(self, trial: "Trial"):
+        config = trial.config.copy()
+
+        config.pop("callbacks", None)  # Remove callbacks
+
+        exclude_results = self._exclude_results.copy()
+
+        # Additional excludes
+        exclude_results += self.excludes
+
+        # Log config keys on each result?
+        if not self.log_config:
+            exclude_results += ["config"]
+
+        # Fill trial ID and name
+        trial_id = trial.trial_id if trial else None
+        trial_name = str(trial) if trial else None
+
+        # Project name for Wandb
+        wandb_project = self.project
+
+        # Grouping
+        wandb_group = self.group or trial.experiment_dir_name if trial else None
+
+        # remove unpickleable items!
+        config = _clean_log(config)
+        config = {
+            key: value for key, value in config.items() if key not in self.excludes
+        }
+
+        wandb_init_kwargs = dict(
+            id=trial_id,
+            name=trial_name,
+            resume=False,
+            reinit=True,
+            allow_val_change=True,
+            group=wandb_group,
+            project=wandb_project,
+            config=config,
+        )
+        wandb_init_kwargs.update(self.kwargs)
+
+        self._start_logging_actor(trial, exclude_results, **wandb_init_kwargs)
+
+    def _start_logging_actor(
+        self, trial: "Trial", exclude_results: List[str], **wandb_init_kwargs
+    ):
+        # Reuse actor if one already exists.
+        # This can happen if the trial is restarted.
+        if trial in self._trial_logging_futures:
+            return
+
+        if not self._remote_logger_class:
+            env_vars = {}
+            # API key env variable is not set if authenticating through `wandb login`
+            if WANDB_ENV_VAR in os.environ:
+                env_vars[WANDB_ENV_VAR] = os.environ[WANDB_ENV_VAR]
+            self._remote_logger_class = ray.remote(
+                num_cpus=0,
+                **_force_on_current_node(),
+                runtime_env={"env_vars": env_vars},
+                max_restarts=-1,
+                max_task_retries=-1,
+            )(self._logger_actor_cls)
+
+        self._trial_queues[trial] = Queue(
+            actor_options={
+                "num_cpus": 0,
+                **_force_on_current_node(),
+                "max_restarts": -1,
+                "max_task_retries": -1,
+            }
+        )
+        self._trial_logging_actors[trial] = self._remote_logger_class.remote(
+            logdir=trial.local_path,
+            queue=self._trial_queues[trial],
+            exclude=exclude_results,
+            to_config=self.AUTO_CONFIG_KEYS,
+            **wandb_init_kwargs,
+        )
+        logging_future = self._trial_logging_actors[trial].run.remote()
+        self._trial_logging_futures[trial] = logging_future
+        self._logging_future_to_trial[logging_future] = trial
+
+    def _signal_logging_actor_stop(self, trial: "Trial"):
+        self._trial_queues[trial].put((_QueueItem.END, None))
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        if trial not in self._trial_logging_actors:
+            self.log_trial_start(trial)
+
+        result = _clean_log(
+            result, video_kwargs=self.video_kwargs, image_kwargs=self.image_kwargs
+        )
+        self._trial_queues[trial].put((_QueueItem.RESULT, result))
+
+    def log_trial_save(self, trial: "Trial"):
+        if self.upload_checkpoints and trial.checkpoint:
+            checkpoint_root = None
+            if isinstance(trial.checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
+                checkpoint_root = trial.checkpoint.path
+
+            if checkpoint_root:
+                self._trial_queues[trial].put((_QueueItem.CHECKPOINT, checkpoint_root))
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        self._signal_logging_actor_stop(trial=trial)
+        self._cleanup_logging_actors()
+
+    def _cleanup_logging_actor(self, trial: "Trial"):
+        del self._trial_queues[trial]
+        del self._trial_logging_futures[trial]
+        ray.kill(self._trial_logging_actors[trial])
+        del self._trial_logging_actors[trial]
+
+    def _cleanup_logging_actors(self, timeout: int = 0, kill_on_timeout: bool = False):
+        """Clean up logging actors that have finished uploading to wandb.
+        Waits for `timeout` seconds to collect finished logging actors.
+
+        Args:
+            timeout: The number of seconds to wait. Defaults to 0 to clean up
+                any immediate logging actors during the run.
+                This is set to a timeout threshold to wait for pending uploads
+                on experiment end.
+            kill_on_timeout: Whether or not to kill and cleanup the logging actor if
+                it hasn't finished within the timeout.
+        """
+
+        futures = list(self._trial_logging_futures.values())
+        done, remaining = ray.wait(futures, num_returns=len(futures), timeout=timeout)
+        for ready_future in done:
+            finished_trial = self._logging_future_to_trial.pop(ready_future)
+            self._cleanup_logging_actor(finished_trial)
+
+        if kill_on_timeout:
+            for remaining_future in remaining:
+                trial = self._logging_future_to_trial.pop(remaining_future)
+                self._cleanup_logging_actor(trial)
+
+    def on_experiment_end(self, trials: List["Trial"], **info):
+        """Wait for the actors to finish their call to `wandb.finish`.
+        This includes uploading all logs + artifacts to wandb."""
+        self._cleanup_logging_actors(timeout=self._upload_timeout, kill_on_timeout=True)
+
+    def __del__(self):
+        if ray.is_initialized():
+            for trial in list(self._trial_logging_actors):
+                self._signal_logging_actor_stop(trial=trial)
+
+            self._cleanup_logging_actors(timeout=2, kill_on_timeout=True)
+
+        self._trial_logging_actors = {}
+        self._trial_logging_futures = {}
+        self._logging_future_to_trial = {}
+        self._trial_queues = {}

python/ray/air/result.py

@@ -0,0 +1,290 @@
+import io
+import json
+import logging
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import pandas as pd
+import pyarrow
+
+import ray
+from ray._private.dict import unflattened_lookup
+from ray.air.constants import (
+    EXPR_ERROR_PICKLE_FILE,
+    EXPR_PROGRESS_FILE,
+    EXPR_RESULT_FILE,
+)
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+@PublicAPI(stability="stable")
+class Result:
+    """The final result of a ML training run or a Tune trial.
+
+    This is the output produced by ``Trainer.fit``.
+    ``Tuner.fit`` outputs a :class:`~ray.tune.ResultGrid` that is a collection
+    of ``Result`` objects.
+
+    This API is the recommended way to access the outputs such as:
+    - checkpoints (``Result.checkpoint``)
+    - the history of reported metrics (``Result.metrics_dataframe``, ``Result.metrics``)
+    - errors encountered during a training run (``Result.error``)
+
+    The constructor is a private API -- use ``Result.from_path`` to create a result
+    object from a directory.
+
+    Attributes:
+        metrics: The latest set of reported metrics.
+        checkpoint: The latest checkpoint.
+        error: The execution error of the Trainable run, if the trial finishes in error.
+        path: Path pointing to the result directory on persistent storage. This can
+            point to a remote storage location (e.g. S3) or to a local location (path
+            on the head node). The path is accessible via the result's associated
+            `filesystem`. For instance, for a result stored in S3 at
+            ``s3://bucket/location``, ``path`` will have the value ``bucket/location``.
+        metrics_dataframe: The full result dataframe of the Trainable.
+            The dataframe is indexed by iterations and contains reported
+            metrics. Note that the dataframe columns are indexed with the
+            *flattened* keys of reported metrics, so the format of this dataframe
+            may be slightly different than ``Result.metrics``, which is an unflattened
+            dict of the latest set of reported metrics.
+        best_checkpoints: A list of tuples of the best checkpoints and
+            their associated metrics. The number of
+            saved checkpoints is determined by :class:`~ray.train.CheckpointConfig`
+            (by default, all checkpoints will be saved).
+    """
+
+    metrics: Optional[Dict[str, Any]]
+    checkpoint: Optional["ray.tune.Checkpoint"]
+    error: Optional[Exception]
+    path: str
+    metrics_dataframe: Optional["pd.DataFrame"] = None
+    best_checkpoints: Optional[
+        List[Tuple["ray.tune.Checkpoint", Dict[str, Any]]]
+    ] = None
+    _storage_filesystem: Optional[pyarrow.fs.FileSystem] = None
+    _items_to_repr = ["error", "metrics", "path", "filesystem", "checkpoint"]
+
+    @property
+    def config(self) -> Optional[Dict[str, Any]]:
+        """The config associated with the result."""
+        if not self.metrics:
+            return None
+        return self.metrics.get("config", None)
+
+    @property
+    def filesystem(self) -> pyarrow.fs.FileSystem:
+        """Return the filesystem that can be used to access the result path.
+
+        Returns:
+            pyarrow.fs.FileSystem implementation.
+        """
+        return self._storage_filesystem or pyarrow.fs.LocalFileSystem()
+
+    def _repr(self, indent: int = 0) -> str:
+        """Construct the representation with specified number of space indent."""
+        from ray.tune.experimental.output import BLACKLISTED_KEYS
+        from ray.tune.result import AUTO_RESULT_KEYS
+
+        shown_attributes = {k: getattr(self, k) for k in self._items_to_repr}
+        if self.error:
+            shown_attributes["error"] = type(self.error).__name__
+        else:
+            shown_attributes.pop("error")
+
+        shown_attributes["filesystem"] = shown_attributes["filesystem"].type_name
+
+        if self.metrics:
+            exclude = set(AUTO_RESULT_KEYS)
+            exclude.update(BLACKLISTED_KEYS)
+            shown_attributes["metrics"] = {
+                k: v for k, v in self.metrics.items() if k not in exclude
+            }
+
+        cls_indent = " " * indent
+        kws_indent = " " * (indent + 2)
+
+        kws = [
+            f"{kws_indent}{key}={value!r}" for key, value in shown_attributes.items()
+        ]
+        kws_repr = ",\n".join(kws)
+        return "{0}{1}(\n{2}\n{0})".format(cls_indent, type(self).__name__, kws_repr)
+
+    def __repr__(self) -> str:
+        return self._repr(indent=0)
+
+    @staticmethod
+    def _read_file_as_str(
+        storage_filesystem: pyarrow.fs.FileSystem,
+        storage_path: str,
+    ) -> str:
+        """Opens a file as an input stream reading all byte content sequentially and
+         decoding read bytes as utf-8 string.
+
+        Args:
+            storage_filesystem: The filesystem to use.
+            storage_path: The source to open for reading.
+
+        Returns:
+            The file contents decoded as a UTF-8 string.
+        """
+
+        with storage_filesystem.open_input_stream(storage_path) as f:
+            return f.readall().decode()
+
+    @classmethod
+    def from_path(
+        cls,
+        path: Union[str, os.PathLike],
+        storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
+    ) -> "Result":
+        """Restore a Result object from local or remote trial directory.
+
+        Args:
+            path: A path of a trial directory on local or remote storage
+                (ex: s3://bucket/path or /tmp/ray_results).
+            storage_filesystem: A custom filesystem to use. If not provided,
+                this will be auto-resolved by pyarrow. If provided, the path
+                is assumed to be prefix-stripped already, and must be a valid path
+                on the filesystem.
+
+        Returns:
+            A :py:class:`Result` object of that trial.
+        """
+        # TODO(justinvyu): Fix circular dependency.
+        from ray.train import Checkpoint
+        from ray.train._internal.storage import (
+            _exists_at_fs_path,
+            _list_at_fs_path,
+            get_fs_and_path,
+        )
+        from ray.train.constants import CHECKPOINT_DIR_NAME
+
+        fs, fs_path = get_fs_and_path(path, storage_filesystem)
+        if not _exists_at_fs_path(fs, fs_path):
+            raise RuntimeError(f"Trial folder {fs_path} doesn't exist!")
+
+        # Restore metrics from result.json
+        result_json_file = Path(fs_path, EXPR_RESULT_FILE).as_posix()
+        progress_csv_file = Path(fs_path, EXPR_PROGRESS_FILE).as_posix()
+        if _exists_at_fs_path(fs, result_json_file):
+            lines = cls._read_file_as_str(fs, result_json_file).split("\n")
+            json_list = [json.loads(line) for line in lines if line]
+            metrics_df = pd.json_normalize(json_list, sep="/")
+            latest_metrics = json_list[-1] if json_list else {}
+        # Fallback to restore from progress.csv
+        elif _exists_at_fs_path(fs, progress_csv_file):
+            metrics_df = pd.read_csv(
+                io.StringIO(cls._read_file_as_str(fs, progress_csv_file))
+            )
+            latest_metrics = (
+                metrics_df.iloc[-1].to_dict() if not metrics_df.empty else {}
+            )
+        else:
+            raise RuntimeError(
+                f"Failed to restore the Result object: Neither {EXPR_RESULT_FILE}"
+                f" nor {EXPR_PROGRESS_FILE} exists in the trial folder!"
+            )
+
+        # Restore all checkpoints from the checkpoint folders
+        checkpoint_dir_names = sorted(
+            _list_at_fs_path(
+                fs,
+                fs_path,
+                file_filter=lambda file_info: file_info.type
+                == pyarrow.fs.FileType.Directory
+                and file_info.base_name.startswith("checkpoint_"),
+            )
+        )
+
+        if checkpoint_dir_names:
+            checkpoints = [
+                Checkpoint(
+                    path=Path(fs_path, checkpoint_dir_name).as_posix(), filesystem=fs
+                )
+                for checkpoint_dir_name in checkpoint_dir_names
+            ]
+
+            metrics = []
+            for checkpoint_dir_name in checkpoint_dir_names:
+                metrics_corresponding_to_checkpoint = metrics_df[
+                    metrics_df[CHECKPOINT_DIR_NAME] == checkpoint_dir_name
+                ]
+                if metrics_corresponding_to_checkpoint.empty:
+                    logger.warning(
+                        "Could not find metrics corresponding to "
+                        f"{checkpoint_dir_name}. These will default to an empty dict."
+                    )
+                metrics.append(
+                    {}
+                    if metrics_corresponding_to_checkpoint.empty
+                    else metrics_corresponding_to_checkpoint.iloc[-1].to_dict()
+                )
+
+            latest_checkpoint = checkpoints[-1]
+            # TODO(justinvyu): These are ordered by checkpoint index, since we don't
+            # know the metric to order these with.
+            best_checkpoints = list(zip(checkpoints, metrics))
+        else:
+            best_checkpoints = latest_checkpoint = None
+
+        # Restore the trial error if it exists
+        error = None
+        error_file_path = Path(fs_path, EXPR_ERROR_PICKLE_FILE).as_posix()
+        if _exists_at_fs_path(fs, error_file_path):
+            with fs.open_input_stream(error_file_path) as f:
+                error = ray.cloudpickle.load(f)
+
+        return Result(
+            metrics=latest_metrics,
+            checkpoint=latest_checkpoint,
+            path=fs_path,
+            _storage_filesystem=fs,
+            metrics_dataframe=metrics_df,
+            best_checkpoints=best_checkpoints,
+            error=error,
+        )
+
+    @PublicAPI(stability="alpha")
+    def get_best_checkpoint(
+        self, metric: str, mode: str
+    ) -> Optional["ray.tune.Checkpoint"]:
+        """Get the best checkpoint from this trial based on a specific metric.
+
+        Any checkpoints without an associated metric value will be filtered out.
+
+        Args:
+            metric: The key for checkpoints to order on.
+            mode: One of ["min", "max"].
+
+        Returns:
+            :class:`Checkpoint <ray.train.Checkpoint>` object, or None if there is
+            no valid checkpoint associated with the metric.
+        """
+        if not self.best_checkpoints:
+            raise RuntimeError("No checkpoint exists in the trial directory!")
+
+        if mode not in ["max", "min"]:
+            raise ValueError(
+                f'Unsupported mode: {mode}. Please choose from ["min", "max"]!'
+            )
+
+        op = max if mode == "max" else min
+        valid_checkpoints = [
+            ckpt_info
+            for ckpt_info in self.best_checkpoints
+            if unflattened_lookup(metric, ckpt_info[1], default=None) is not None
+        ]
+
+        if not valid_checkpoints:
+            raise RuntimeError(
+                f"Invalid metric name {metric}! "
+                f"You may choose from the following metrics: {self.metrics.keys()}."
+            )
+
+        return op(valid_checkpoints, key=lambda x: unflattened_lookup(metric, x[1]))[0]

python/ray/air/session.py

@@ -0,0 +1 @@
+from ray.train._internal.session import *  # noqa: F401,F403

python/ray/air/tests/_test_experiment_restore_run.py

@@ -0,0 +1,180 @@
+import collections
+import json
+import os
+import random
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import ray
+from ray import train, tune
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.tune.experiment import Trial
+
+RUNNER_TYPE = os.environ.get("RUNNER_TYPE", "trainer")
+STORAGE_PATH = os.environ.get("STORAGE_PATH", "/tmp/ray_results")
+EXP_NAME = os.environ.get("EXP_NAME", "restore_integration_test")
+CALLBACK_DUMP_FILE = os.environ.get(
+    "CALLBACK_DUMP_FILE", "/tmp/callback_dump_file.json"
+)
+CSV_DATA_FILE = os.environ.get("CSV_DATA_FILE", "/tmp/dummy.csv")
+
+TIME_PER_ITER_S = float(os.environ.get("TIME_PER_ITER_S", "0.5"))
+NUM_TRIALS = int(os.environ.get("NUM_TRIALS", "1"))
+MAX_CONCURRENT_TRIALS = int(os.environ.get("MAX_CONCURRENT_TRIALS", "2"))
+ITERATIONS_PER_TRIAL = int(os.environ.get("ITERATIONS_PER_TRIAL", "64"))
+
+
+class StatefulCallback(tune.Callback):
+    def __init__(self):
+        self._trial_iterations = collections.defaultdict(list)
+
+    def on_trial_result(
+        self,
+        iteration: int,
+        trials: List["Trial"],
+        trial: "Trial",
+        result: Dict,
+        **info,
+    ):
+        self._trial_iterations[trial.trial_id].append(result["training_iteration"])
+
+    def on_experiment_end(self, trials: List["Trial"], **info):
+        # Save callback contents to file
+        with open(CALLBACK_DUMP_FILE, "w") as f:
+            json.dump(self.get_state(), f, indent=2)
+
+    def get_state(self) -> Optional[Dict]:
+        return {"trial_iters": self._trial_iterations.copy()}
+
+    def set_state(self, state: Dict):
+        self._trial_iterations = state["trial_iters"]
+
+
+class StatefulSearcher(tune.search.Searcher):
+    def __init__(
+        self,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+    ):
+        super().__init__(metric=metric, mode=mode)
+        self._trial_count = 0
+
+    def suggest(self, trial_id: str) -> Optional[Dict]:
+        self._trial_count += 1
+        return {"id": self._trial_count}
+
+    def on_trial_complete(
+        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
+    ) -> None:
+        pass
+
+    def save(self, checkpoint_path: str):
+        with open(checkpoint_path, "w") as f:
+            json.dump({"trial_count": self._trial_count}, f)
+
+    def restore(self, checkpoint_path: str):
+        with open(checkpoint_path, "r") as f:
+            state = json.load(f)
+        self._trial_count = state["trial_count"]
+
+
+def train_fn(config: dict, data: Optional[dict] = None):
+    checkpoint = train.get_checkpoint()
+    start = load_dict_checkpoint(checkpoint)["iteration"] + 1 if checkpoint else 1
+
+    training_started_marker = Path(
+        os.environ.get("RUN_STARTED_MARKER", "/tmp/does-not-exist")
+    )
+    if training_started_marker.exists():
+        # Multiple workers may be trying to delete the same marker
+        try:
+            training_started_marker.unlink()
+        except FileNotFoundError:
+            pass
+
+    for iteration in range(start, ITERATIONS_PER_TRIAL + 1):
+        time.sleep(TIME_PER_ITER_S)
+
+        with create_dict_checkpoint({"iteration": iteration}) as checkpoint:
+            train.report({"score": random.random()}, checkpoint=checkpoint)
+
+
+def tuner(experiment_path: str, run_config: tune.RunConfig) -> tune.ResultGrid:
+    trainable = tune.with_resources(train_fn, resources={"CPU": 1})
+    trainable = tune.with_parameters(trainable, data={"dummy_data": [1, 2, 3]})
+
+    if tune.Tuner.can_restore(experiment_path):
+        tuner = tune.Tuner.restore(
+            experiment_path, trainable=trainable, resume_errored=True
+        )
+    else:
+        tuner = tune.Tuner(
+            trainable,
+            run_config=run_config,
+            tune_config=tune.TuneConfig(
+                num_samples=8,
+                max_concurrent_trials=2,
+                search_alg=StatefulSearcher(),
+            ),
+        )
+
+    result_grid = tuner.fit()
+    return result_grid
+
+
+def trainer(experiment_path: str, run_config: train.RunConfig) -> train.Result:
+    dataset_size = 128
+    num_workers = 4
+
+    def train_loop_per_worker(config):
+        # Wrap the other train_fn with a check for the dataset.
+        assert train.get_dataset_shard("train")
+        train_fn(config)
+
+    datasets = {
+        "train": ray.data.range(dataset_size),
+        "valid": ray.data.read_csv(CSV_DATA_FILE),
+    }
+
+    if DataParallelTrainer.can_restore(experiment_path):
+        trainer = DataParallelTrainer.restore(
+            experiment_path,
+            datasets=datasets,
+            train_loop_per_worker=train_loop_per_worker,
+        )
+    else:
+        trainer = DataParallelTrainer(
+            train_loop_per_worker,
+            datasets=datasets,
+            scaling_config=train.ScalingConfig(
+                num_workers=num_workers, trainer_resources={"CPU": 0}
+            ),
+            run_config=run_config,
+        )
+
+    result = trainer.fit()
+    return result
+
+
+if __name__ == "__main__":
+    experiment_path = os.path.join(STORAGE_PATH, EXP_NAME)
+
+    ray.init()
+
+    run_config = train.RunConfig(
+        storage_path=STORAGE_PATH,
+        name=EXP_NAME,
+        checkpoint_config=train.CheckpointConfig(num_to_keep=1),
+        callbacks=[StatefulCallback()],
+    )
+
+    if RUNNER_TYPE == "tuner":
+        tuner(experiment_path, run_config)
+    elif RUNNER_TYPE == "trainer":
+        trainer(experiment_path, run_config)
+    else:
+        raise NotImplementedError(
+            "`RUNNER_TYPE` environment var must be one of ['tuner', 'trainer']"
+        )

python/ray/air/tests/conftest.py

@@ -0,0 +1,23 @@
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+import copy
+
+import pytest
+
+import ray
+from ray.tests.conftest import pytest_runtest_makereport  # noqa
+
+
+@pytest.fixture
+def restore_data_context(request):
+    """Restore any DataContext changes after the test runs"""
+    original = copy.deepcopy(ray.data.context.DataContext.get_current())
+    yield
+    ray.data.context.DataContext._set_current(original)
+
+
+@pytest.fixture
+def disable_fallback_to_object_extension(request, restore_data_context):
+    """Disables fallback to ArrowPythonObjectType"""
+    ray.data.context.DataContext.get_current().enable_fallback_to_arrow_object_ext_type = (
+        False
+    )

python/ray/air/tests/execution/test_barrier.py

@@ -0,0 +1,54 @@
+from typing import Optional, Type
+
+import pytest
+
+from ray.air.execution._internal.barrier import Barrier
+
+
+def _raise(exception_type: Type[Exception] = RuntimeError, msg: Optional[str] = None):
+    def _raise_exception(*args, **kwargs):
+        raise exception_type(msg)
+
+    return _raise_exception
+
+
+def test_barrier_max_results():
+    """Test the `max_results` attribute.
+
+    - Set max_results=10
+    - Assert that the barrier completion callback is not invoked with num_results<10
+    - Assert that callback is invoked with num_results=10
+    - Assert that callback is not invoked again when more events arrive
+    - Assert that more events can arrive without triggering the callback after resetting
+    """
+    barrier = Barrier(max_results=10, on_completion=_raise(AssertionError))
+
+    for i in range(9):
+        barrier.arrive(i)
+
+    assert not barrier.completed
+
+    # Will trigger the on_completion callback
+    with pytest.raises(AssertionError):
+        barrier.arrive(10)
+
+    assert barrier.completed
+
+    assert barrier.num_results == 10
+
+    # Further events will not trigger callback again
+    barrier.arrive(11)
+
+    barrier.reset()
+
+    assert not barrier.completed
+
+    # After flushing more events can arrive
+    barrier.arrive(12)
+    assert barrier.num_results == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_e2e_train_flow.py

@@ -0,0 +1,266 @@
+import random
+from typing import Any, List, Optional
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal import Barrier
+from ray.air.execution._internal.actor_manager import RayActorManager
+from ray.air.execution._internal.tracked_actor import TrackedActor
+from ray.exceptions import RayActorError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+class Actor:
+    """Simple actor for testing an execution flow.
+
+    This actor can fail in these ways:
+
+    1. On init if ``actor_init_kill`` is passed as a kwarg
+    2. On setup_1() if ``actor_setup_kill`` is passed as a kwarg (RayActorError)
+    3. On setup_1() if ``actor_setup_fail`` is passed as a kwarg (RayTaskError)
+    4. On train() if ``actor_train_kill`` is passed as a kwarg (RayTaskError)
+    5. On train() if ``actor_train_fail`` is passed as a kwarg (RayTaskError)
+    """
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+        if self.kwargs.get("actor_init_kill"):
+            raise RuntimeError("INIT")
+
+    def get_kwargs(self):
+        return self.kwargs
+
+    def setup_1(self):
+        if self.kwargs.get("actor_setup_kill"):
+            raise SystemExit
+
+        if self.kwargs.get("actor_setup_fail"):
+            raise RuntimeError("Setup")
+
+        return True
+
+    def setup_2(self):
+        return True
+
+    def train(self, value: float) -> float:
+        if value == 4:
+            if self.kwargs.get("actor_train_kill"):
+                # SystemExit will invoke a RayActorError
+                raise SystemExit
+
+            if self.kwargs.get("actor_train_fail"):
+                # RuntimeError will invoke a RayTaskError
+                raise RuntimeError("TASK")
+
+        return value
+
+
+class TrainFlow:
+    """This is a Ray Train-like execution flow.
+
+    - We want to run 4 actors in total ("trials")
+    - Each actor runs two init functions
+    - We train all actors in parallel for 10 iterations
+    - Errors can come up on actor construction, in the init functions,
+      or during training
+    - When an actor fails, restart that actor
+    - When a task fails, stop actor, and restart
+    """
+
+    def __init__(
+        self, actor_manager: RayActorManager, errors: Optional[List[str]] = None
+    ):
+        self._actor_manager = actor_manager
+        self._finished = False
+
+        self._actors_to_run = 4
+        self._tracked_actors = []
+        self._actors_stopped = 0
+
+        self._actors_to_replace = set()
+
+        self._ready_actors = set()
+        self._training_barrier = Barrier(
+            max_results=self._actors_to_run,
+            on_completion=self.training_barrier_completed,
+        )
+        self._restart_training = None
+
+        self._training_iter = 0
+        self._results = []
+
+        self._errors = errors
+
+    def setup_actors(self):
+        for actor_id in range(self._actors_to_run):
+            error_kwargs = {}
+            if self._errors:
+                error = random.choice(self._errors)
+                error_kwargs[error] = True
+
+            print("Actor", actor_id, "will be failing with", error_kwargs)
+
+            tracked_actor = self._actor_manager.add_actor(
+                cls=Actor,
+                kwargs={"id": actor_id, **error_kwargs},
+                resource_request=ResourceRequest([{"CPU": 1}]),
+                on_start=self.actor_started,
+                on_stop=self.actor_stopped,
+                on_error=self.actor_error,
+            )
+            self._tracked_actors.append(tracked_actor)
+
+    def actor_started(self, tracked_actor: TrackedActor):
+        self._actor_manager.schedule_actor_task(
+            tracked_actor,
+            "setup_1",
+            on_error=self.setup_error,
+            on_result=self.setup_1_result,
+        )
+
+    def actor_stopped(self, tracked_actor: TrackedActor):
+        self._ready_actors.discard(tracked_actor)
+
+        if tracked_actor in self._actors_to_replace:
+            self._replace_actor(tracked_actor=tracked_actor)
+        else:
+            self._actors_stopped += 1
+            self._finished = self._actors_stopped >= self._actors_to_run
+
+    def actor_error(self, tracked_actor: TrackedActor, exception: Exception):
+        self._ready_actors.discard(tracked_actor)
+        self._replace_actor(tracked_actor=tracked_actor)
+
+    def _replace_actor(self, tracked_actor: TrackedActor):
+        actor_index = self._tracked_actors.index(tracked_actor)
+
+        replacement_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={"id": actor_index},
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+
+        self._tracked_actors[actor_index] = replacement_actor
+
+    def setup_1_result(self, tracked_actor: TrackedActor, result: Any):
+        self._actor_manager.schedule_actor_task(
+            tracked_actor,
+            "setup_2",
+            on_error=self.setup_error,
+            on_result=self.setup_2_result,
+        )
+
+    def setup_2_result(self, tracked_actor: TrackedActor, result: Any):
+        self._ready_actors.add(tracked_actor)
+
+        if len(self._ready_actors) == self._actors_to_run:
+            self.continue_training()
+
+    def setup_error(self, tracked_actor: TrackedActor, exception: Exception):
+        if isinstance(exception, RayActorError):
+            return
+
+        self._actors_to_replace.add(tracked_actor)
+        self._actor_manager.remove_actor(tracked_actor)
+
+    def continue_training(self):
+        if self._restart_training:
+            self._training_iter = self._restart_training
+        else:
+            self._training_iter += 1
+
+        self._training_barrier.reset()
+        self._actor_manager.schedule_actor_tasks(
+            self._tracked_actors,
+            "train",
+            args=(self._training_iter,),
+            on_result=self._training_barrier.arrive,
+            on_error=self.training_error,
+        )
+
+    def training_barrier_completed(self, barrier: Barrier):
+        self._results.append([res for _, res in barrier.get_results()])
+        self._restart_training = None
+
+        # If less than 10 epochs, continue training
+        if self._training_iter < 10:
+            return self.continue_training()
+
+        # Else, training finished
+        for tracked_actor in self._tracked_actors:
+            self._actor_manager.remove_actor(tracked_actor)
+
+    def training_error(self, tracked_actor: TrackedActor, exception: Exception):
+        self._restart_training = self._training_iter
+
+        if isinstance(exception, RayActorError):
+            return
+
+        self._actors_to_replace.add(tracked_actor)
+        self._ready_actors.discard(tracked_actor)
+        self._actor_manager.remove_actor(tracked_actor)
+
+    def run(self):
+        self.setup_actors()
+
+        while not self._finished:
+            self._actor_manager.next()
+
+    def get_results(self) -> List[List[float]]:
+        return self._results
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize(
+    "errors",
+    [
+        None,
+        "actor_init_kill",
+        "actor_setup_kill",
+        "actor_setup_fail",
+        "actor_train_kill",
+        "actor_train_fail",
+        # Chaos - every actor fails somehow, but in different ways
+        [
+            "actor_init_kill",
+            "actor_setup_kill",
+            "actor_setup_fail",
+            "actor_train_kill",
+            "actor_train_fail",
+        ],
+    ],
+)
+def test_e2e(ray_start_4_cpus, resource_manager_cls, errors):
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    if errors and isinstance(errors, str):
+        errors = [errors]
+
+    flow = TrainFlow(actor_manager=actor_manager, errors=errors)
+    flow.run()
+
+    results = flow.get_results()
+
+    assert results == [[i] * 4 for i in range(1, 11)], results
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_e2e_tune_flow.py

@@ -0,0 +1,227 @@
+import random
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal.actor_manager import RayActorManager
+from ray.air.execution._internal.tracked_actor import TrackedActor
+from ray.exceptions import RayActorError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+class Actor:
+    """Simple actor for testing an execution flow.
+
+    This actor can fail in three ways:
+
+    1. On init if ``actor_error_init`` is passed as a kwarg
+    2. On run() if ``actor_error_task`` is passed as a kwarg (RayActorError)
+    3. On run() if ``task_error`` is passed as a kwarg (RayTaskError)
+    """
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+        if self.kwargs.get("actor_error_init"):
+            raise RuntimeError("INIT")
+
+    def get_kwargs(self):
+        return self.kwargs
+
+    def run(self, value: float) -> float:
+        if value == 2:
+            if self.kwargs.get("actor_error_task"):
+                # SystemExit will invoke a RayActorError
+                raise SystemExit
+
+            if self.kwargs.get("task_error"):
+                # RuntimeError will invoke a RayTaskError
+                raise RuntimeError("TASK")
+
+        return value
+
+
+class TuneFlow:
+    """This is a Ray Tune-like execution flow.
+
+    - We want to run 10 actors in total ("trials")
+    - Each actor collects 11 results sequentially
+    - We schedule up to 6 actors at the same time
+    - Every step, we see if we should add any new actors
+    - Otherwise, we just yield control to the event manager and process events one
+      by one
+    - When an actor is started, start training flow
+    - When a result comes in, schedule next future
+      - If this is the 11th result, stop actor
+    - When the last actor is stopped, set state to finished
+
+    - When an actor fails, restart
+    - When a task fails, stop actor, and restart
+    """
+
+    def __init__(
+        self, actor_manager: RayActorManager, errors: Optional[List[str]] = None
+    ):
+        self._actor_manager = actor_manager
+        self._finished = False
+
+        self._actors_to_run = 10
+        self._actors_started = 0
+        self._actors_stopped = 0
+        self._max_pending = 6
+
+        self._actor_to_id = {}
+        self._results = defaultdict(list)
+
+        self._errors = errors
+
+    def maybe_add_actors(self):
+        if self._actors_started >= self._actors_to_run:
+            return
+
+        if self._actor_manager.num_pending_actors >= self._max_pending:
+            return
+
+        error_kwargs = {}
+        if self._errors:
+            error = random.choice(self._errors)
+            error_kwargs[error] = True
+
+        actor_id = self._actors_started
+
+        print("Actor", actor_id, "will be failing with", error_kwargs)
+
+        tracked_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={"id": actor_id, **error_kwargs},
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+        self._actor_to_id[tracked_actor] = actor_id
+
+        self._actors_started += 1
+
+    def actor_started(self, tracked_actor: TrackedActor):
+        self._actor_manager.schedule_actor_task(
+            tracked_actor,
+            "run",
+            kwargs={"value": 0},
+            on_error=self.task_error,
+            on_result=self.task_result,
+        )
+
+    def actor_stopped(self, tracked_actor: TrackedActor):
+        self._actors_stopped += 1
+        self._finished = self._actors_stopped >= self._actors_to_run
+
+    def actor_error(self, tracked_actor: TrackedActor, exception: Exception):
+        actor_id = self._actor_to_id.pop(tracked_actor)
+
+        replacement_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={
+                "id": actor_id,
+                "actor_error_init": False,
+                "actor_error_task": False,
+                "task_error": False,
+            },
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+
+        self._actor_to_id[replacement_actor] = actor_id
+
+    def task_result(self, tracked_actor: TrackedActor, result: float):
+        actor_id = self._actor_to_id[tracked_actor]
+        self._results[actor_id].append(result)
+
+        if result == 10:
+            self._actor_manager.remove_actor(tracked_actor)
+        else:
+            self._actor_manager.schedule_actor_task(
+                tracked_actor,
+                "run",
+                kwargs={"value": result + 1},
+                on_result=self.task_result,
+                on_error=self.task_error,
+            )
+
+    def task_error(self, tracked_actor: TrackedActor, exception: Exception):
+        if isinstance(exception, RayActorError):
+            return
+
+        self._actors_stopped -= 1  # account for extra stop
+        self._actor_manager.remove_actor(tracked_actor)
+        actor_id = self._actor_to_id.pop(tracked_actor)
+
+        replacement_actor = self._actor_manager.add_actor(
+            cls=Actor,
+            kwargs={
+                "id": actor_id,
+                "actor_error_init": False,
+                "actor_error_task": False,
+                "task_error": False,
+            },
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=self.actor_started,
+            on_stop=self.actor_stopped,
+            on_error=self.actor_error,
+        )
+        self._actor_to_id[replacement_actor] = actor_id
+
+    def run(self):
+        while not self._finished:
+            self.maybe_add_actors()
+            self._actor_manager.next(timeout=1)
+
+    def get_results(self) -> Dict[int, List[float]]:
+        return self._results
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize(
+    "errors",
+    [
+        None,
+        "actor_error_init",
+        "actor_error_task",
+        "task_error",
+        # Chaos - every actor fails somehow, but in different ways
+        ["actor_error_init", "actor_error_task", "task_error"],
+    ],
+)
+def test_e2e(ray_start_4_cpus, resource_manager_cls, errors):
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    if errors and isinstance(errors, str):
+        errors = [errors]
+
+    flow = TuneFlow(actor_manager=actor_manager, errors=errors)
+    flow.run()
+
+    results = flow.get_results()
+
+    assert all(res[-1] == 10 for res in results.values()), results
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_event_manager.py

@@ -0,0 +1,224 @@
+import time
+from typing import Any, Type
+
+import pytest
+
+import ray
+from ray.air.execution._internal import Barrier
+from ray.air.execution._internal.event_manager import RayEventManager
+from ray.exceptions import RayTaskError
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+def succeeding(ret: Any = None) -> Any:
+    return ret
+
+
+@ray.remote
+def failing(exc: Type[Exception], *args) -> None:
+    raise exc(*args)
+
+
+@ray.remote
+def sleeping(seconds: int, result: Any) -> Any:
+    time.sleep(seconds)
+    return result
+
+
+def test_track_future_success(ray_start_4_cpus):
+    """Schedule a future that return successfully.
+
+    Check that the on_result callback was triggered.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_result(result: Any):
+        seen.add(result)
+
+    event_manager.track_future(succeeding.remote("a"), on_result=on_result)
+
+    event_manager.wait()
+    assert "a" in seen
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_success_no_callback(ray_start_4_cpus):
+    """Schedule a future that return successfully.
+
+    Check that passing no callback still succeeds.
+    """
+    event_manager = RayEventManager()
+
+    event_manager.track_future(succeeding.remote("a"))
+
+    event_manager.wait()
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_error(ray_start_4_cpus):
+    """Schedule a future that fails.
+
+    Check that the on_error callback was triggered.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    class CustomError(RuntimeError):
+        pass
+
+    def on_error(exception: Exception):
+        seen.add(exception)
+
+    event_manager.track_future(failing.remote(CustomError), on_error=on_error)
+
+    event_manager.wait()
+    assert isinstance(seen.pop(), CustomError)
+
+    assert not event_manager._tracked_futures
+
+
+def test_track_future_error_no_callback(ray_start_4_cpus):
+    """Schedule a future that fails.
+
+    Check that passing no callback raises the original error.
+    """
+    event_manager = RayEventManager()
+
+    event_manager.track_future(failing.remote(RuntimeError))
+
+    with pytest.raises(RuntimeError):
+        event_manager.wait()
+
+    assert not event_manager._tracked_futures
+
+
+@pytest.mark.parametrize("results_per_wait", [None, 1, 5, 10, 100])
+def test_many_futures(ray_start_4_cpus, results_per_wait):
+    """Schedule 500 succeeding and failing futures.
+
+    Check that the callbacks get triggered correctly, independent of the number
+    of results we await per call to RayEventManager.wait().
+    """
+    num_futures = 500
+
+    event_manager = RayEventManager()
+
+    seen_results = set()
+    seen_errors = set()
+
+    def on_result(result: Any):
+        seen_results.add(result)
+
+    def on_error(exception: RayTaskError):
+        seen_errors.add(exception.cause.args[0])
+
+    for i in range(num_futures):
+        event_manager.track_futures(
+            [
+                succeeding.remote("a" + str(i)),
+                failing.remote(RuntimeError, "b" + str(i)),
+            ],
+            on_result=on_result,
+            on_error=on_error,
+        )
+
+    while event_manager.num_futures > 0:
+        event_manager.wait(num_results=results_per_wait)
+
+    for i in range(num_futures):
+        assert "a" + str(i) in seen_results
+        assert "b" + str(i) in seen_errors
+
+
+def test_timeout(ray_start_4_cpus):
+    """Test the timeout parameter.
+
+    Start 4 tasks: Two succeed immediately, two after 1 second.
+
+    After waiting for 0.5 seconds, the first two tasks should have returned.
+    After waiting for up to 5 seconds, the other two tasks should have returned.
+    But because the tasks take only 0.5 seconds to run, we should have waited
+    way less than 5 seconds.
+    """
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_result(result: Any):
+        seen.add(result)
+
+    event_manager.track_futures(
+        [
+            succeeding.remote("a"),
+            succeeding.remote("b"),
+            sleeping.remote(1, "c"),
+            sleeping.remote(1, "d"),
+        ],
+        on_result=on_result,
+    )
+
+    start = time.monotonic()
+    event_manager.wait(num_results=None, timeout=0.5)
+    assert "a" in seen
+    assert "b" in seen
+    assert "c" not in seen
+    assert "d" not in seen
+
+    event_manager.wait(num_results=None, timeout=5)
+    taken = time.monotonic() - start
+
+    assert "c" in seen
+    assert "d" in seen
+
+    # Should have returned much earlier than after 5 seconds
+    assert taken < 3
+
+    assert not event_manager._tracked_futures
+
+
+def test_task_barrier(ray_start_4_cpus):
+    event_manager = RayEventManager()
+
+    seen = set()
+
+    def on_completion(barrier: Barrier):
+        seen.update(barrier.get_results())
+
+    barrier = Barrier(max_results=4, on_completion=on_completion)
+
+    event_manager.track_futures(
+        [
+            succeeding.remote("a"),
+            succeeding.remote("b"),
+            succeeding.remote("c"),
+            succeeding.remote("d"),
+            sleeping.remote(2, "e"),
+        ],
+        on_result=barrier.arrive,
+    )
+
+    event_manager.wait(num_results=4)
+
+    assert "a" in seen
+    assert "b" in seen
+    assert "c" in seen
+    assert "d" in seen
+    assert "e" not in seen
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_resource_manager_fixed.py

@@ -0,0 +1,178 @@
+import pytest
+
+import ray
+from ray.air.execution.resources.fixed import FixedResourceManager
+from ray.air.execution.resources.request import ResourceRequest
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+REQUEST_2_CPU = ResourceRequest([{"CPU": 2}])
+REQUEST_4_CPU = ResourceRequest([{"CPU": 4}])
+REQUEST_1_2_CPU = ResourceRequest([{"CPU": 1}, {"CPU": 2}])
+REQUEST_0_2_CPU = ResourceRequest([{"CPU": 0}, {"CPU": 2}])
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def test_acquire_return_resources(ray_start_4_cpus):
+    manager = FixedResourceManager(total_resources={"CPU": 4})
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_4_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_4_CPU)
+
+    assert manager.has_resources_ready(REQUEST_4_CPU)
+
+    ready_2 = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_4_CPU)
+
+    manager.free_resources(ready_2)
+
+    assert manager.has_resources_ready(REQUEST_4_CPU)
+
+
+def test_numerical_error(ray_start_4_cpus):
+    """Make sure we don't run into numerical errors when using fractional resources.
+
+    Legacy test: test_trial_runner::TrialRunnerTest::testResourceNumericalError
+    """
+    manager = FixedResourceManager(
+        total_resources={"CPU": 0.99, "GPU": 0.99, "a": 0.99}
+    )
+    resource_request = ResourceRequest([{"CPU": 0.33, "GPU": 0.33, "a": 0.33}])
+
+    for i in range(3):
+        manager.request_resources(resource_request)
+        assert manager.acquire_resources(
+            resource_request=resource_request
+        ), manager._available_resources
+
+    assert manager._available_resources["CPU"] == 0
+    assert manager._available_resources["GPU"] == 0
+    assert manager._available_resources["a"] == 0
+
+
+def test_bind_two_bundles(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works.
+
+    - Request resources with 2 bundles (1 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 1 CPU and 2 CPUs
+    """
+    manager = FixedResourceManager()
+    manager.request_resources(REQUEST_1_2_CPU)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 1
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 1
+    assert sum(v for k, v in res2.items() if k.startswith("CPU")) == 2
+
+
+def test_bind_empty_head_bundle(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works with empty head.
+
+    - Request resources with 2 bundles (0 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 0 CPU and 2 CPUs
+    """
+    manager = FixedResourceManager()
+    assert REQUEST_0_2_CPU.head_bundle_is_empty
+    manager.request_resources(REQUEST_0_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_0_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_0_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 0
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 0
+    assert sum(v for k, v in res2.items() if k.startswith("CPU")) == 2
+
+
+@pytest.mark.parametrize("strategy", ["STRICT_PACK", "PACK", "SPREAD", "STRICT_SPREAD"])
+def test_strategy(ray_start_4_cpus, strategy):
+    """The fixed resoure manager does not support STRICT placement strategies."""
+    manager = FixedResourceManager()
+
+    req = ResourceRequest([{"CPU": 2}], strategy=strategy)
+
+    if strategy.startswith("STRICT_"):
+        with pytest.raises(RuntimeError):
+            manager.request_resources(req)
+    else:
+        manager.request_resources(req)
+
+
+@pytest.mark.parametrize("strategy", ["STRICT_PACK", "PACK", "SPREAD", "STRICT_SPREAD"])
+def test_strategy_nested(ray_start_4_cpus, strategy):
+    """The fixed resoure manager does not support STRICT_SPREAD within a PG."""
+
+    @ray.remote
+    def nested_test():
+        manager = FixedResourceManager()
+
+        req = ResourceRequest([{"CPU": 2}], strategy=strategy)
+
+        if strategy == "STRICT_SPREAD":
+            with pytest.raises(RuntimeError):
+                manager.request_resources(req)
+        else:
+            manager.request_resources(req)
+
+    pg = ray.util.placement_group([{"CPU": 2}])
+    ray.wait([pg.ready()])
+
+    try:
+        ray.get(
+            nested_test.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=pg, placement_group_capture_child_tasks=True
+                )
+            ).remote()
+        )
+    finally:
+        ray.util.remove_placement_group(pg)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_resource_manager_placement_group.py

@@ -0,0 +1,409 @@
+import time
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air.execution.resources.placement_group import PlacementGroupResourceManager
+from ray.air.execution.resources.request import ResourceRequest
+
+REQUEST_2_CPU = ResourceRequest([{"CPU": 2}])
+REQUEST_1_2_CPU = ResourceRequest([{"CPU": 1}, {"CPU": 2}])
+REQUEST_0_2_CPU = ResourceRequest([{"CPU": 0}, {"CPU": 2}])
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def _count_pg_states():
+    counter = Counter()
+    for _, pg_info in ray.util.placement_group_table().items():
+        counter[pg_info["state"]] += 1
+    return counter
+
+
+def test_request_cancel_resources(ray_start_4_cpus):
+    """Test that canceling a resource request clears the PG futures.
+
+    - Create request
+    - Assert actual PG is created
+    - Cancel request
+    - Assert staging future is removed
+    - Assert actual PG is removed
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+
+    # Could be pending or created
+    pg_states = _count_pg_states()
+    assert pg_states["PENDING"] + pg_states["CREATED"] == 1
+    assert pg_states["REMOVED"] == 0
+
+    assert manager.get_resource_futures()
+
+    manager.cancel_resource_request(REQUEST_2_CPU)
+
+    assert not manager.get_resource_futures()
+
+    pg_states = _count_pg_states()
+    assert pg_states["PENDING"] + pg_states["CREATED"] == 0
+    assert pg_states["REMOVED"] == 1
+
+
+def test_acquire_return_resources(ray_start_4_cpus):
+    """Tests that acquiring and returning resources works.
+
+    - At the start, no resources should be ready (no PG scheduled)
+    - Request resources for 2 CPUs
+    - (wait until they are ready)
+    - Assert that these 2 CPUs are available to be acquired
+    - Acquire
+    - Assert that there are no 2 CPU resources available anymore
+    - Free resources
+    - Assert that the 2 CPU resources are still not available (no new request)
+        - This is also tested in includes test_request_cancel_resources
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    # Request PG
+    manager.request_resources(REQUEST_2_CPU)
+
+    # Wait until ready
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+
+    # PG exists
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 1
+    assert pg_states["REMOVED"] == 0
+
+    # Acquire PG
+    acquired = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    # Free resources
+    manager.free_resources(acquired)
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    # PG still exists
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 0
+    assert pg_states["REMOVED"] == 1
+
+
+def test_request_pending(ray_start_4_cpus):
+    """Test that requesting too many resources leads to pending PGs.
+
+    - Cluster of 4 CPUs
+    - Request 3 PGs a 2 CPUs
+    - Acquire 2 PGs
+    - Assert no resources are available anymore
+    - Return both PGs
+    - Assert resources are available again
+    - Cancel request
+    - Assert no resources are available again
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_2_CPU)
+
+    # Wait until some are ready
+    ray.wait(manager.get_resource_futures(), num_returns=2)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+    assert len(manager.get_resource_futures()) == 1
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 2
+    assert pg_states["PENDING"] == 1
+    assert pg_states["REMOVED"] == 0
+
+    acq1 = manager.acquire_resources(REQUEST_2_CPU)
+    acq2 = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    manager.free_resources(acq1)
+    manager.free_resources(acq2)
+
+    # Third PG becomes ready
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 1
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 2
+
+    manager.cancel_resource_request(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 0
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 3
+
+
+def test_acquire_unavailable(ray_start_4_cpus):
+    """Test that acquiring resources that are not available returns None.
+
+    - Try to acquire
+    - Assert this does not work
+    - Request resources
+    - Wait until ready
+    - Acquire
+    - Assert this did work
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert not manager.acquire_resources(REQUEST_2_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+    assert manager.acquire_resources(REQUEST_2_CPU)
+
+
+def test_bind_two_bundles(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works.
+
+    - Request PG with 2 bundles (1 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 1 CPU and 2 CPUs
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    manager.request_resources(REQUEST_1_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert res1 == {"CPU": 1}
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert res1 == {"CPU": 1}
+    assert res2 == {"CPU": 2}
+
+
+def test_bind_empty_head_bundle(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works with empty head.
+
+    - Request PG with 2 bundles (0 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 0 CPU and 2 CPUs
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    assert REQUEST_0_2_CPU.head_bundle_is_empty
+    manager.request_resources(REQUEST_0_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_0_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_0_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert res1 == {}
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert res1 == {}
+    assert res2 == {"CPU": 2}
+
+
+def test_capture_child_tasks(ray_start_4_cpus):
+    """Test that child tasks are captured when creating placement groups.
+
+    - Request PG with 2 bundles (1 CPU and 2 CPUs)
+    - Bind a remote task that needs 2 CPUs to run
+    - Assert that it can be scheduled from within the first bundle
+
+    This is only the case if child tasks are captured in the placement groups, as
+    there is only 1 CPU available outside (on a 4 CPU cluster). The 2 CPUs
+    thus have to come from the placement group.
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    manager.request_resources(REQUEST_1_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def needs_cpus():
+        return "Ok"
+
+    @ray.remote
+    def spawn_child_task(num_cpus: int):
+        return ray.get(needs_cpus.options(num_cpus=num_cpus).remote())
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([spawn_child_task])
+
+    res = ray.get(av1.remote(2), timeout=2.0)
+
+    assert res
+
+
+def test_clear_state(ray_start_4_cpus):
+    """Test that clearing state will remove existing placement groups.
+
+    - Create resource request
+    - Wait until PG is scheduled
+    - Assert that Ray PG is created
+    - Call `mgr.clear()`
+    - Assert that resources are not ready anymore
+    - Assert that Ray PG is removed
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+    manager.request_resources(REQUEST_1_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 1
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 0
+
+    manager.clear()
+
+    assert not manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    pg_states = _count_pg_states()
+    assert pg_states["CREATED"] == 0
+    assert pg_states["PENDING"] == 0
+    assert pg_states["REMOVED"] == 1
+
+
+def test_internal_state(ray_start_4_cpus):
+    """Test internal state mappings of the placement group manager.
+
+    This test makes assumptions and assertions around the internal state transition
+    of private properties of the placement group resource manager.
+
+    If you change internal handling logic of the manager, you may need to change this
+    test as well.
+    """
+    manager = PlacementGroupResourceManager(update_interval_s=0)
+
+    assert manager.update_interval_s == 0
+
+    manager.has_resources_ready(REQUEST_2_CPU)
+
+    # The key may exist but the set should be empty
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+
+    ####
+    # 1. Request, wait until ready, cancel
+
+    # Request resources
+    manager.request_resources(REQUEST_2_CPU)
+
+    # PG should be staged
+    assert manager._request_to_staged_pgs[REQUEST_2_CPU]
+    pg = list(manager._request_to_staged_pgs[REQUEST_2_CPU])[0]
+    assert manager._pg_to_request[pg] == REQUEST_2_CPU
+
+    # Staging future should exist
+    assert manager._pg_to_staging_future[pg]
+    fut = manager._pg_to_staging_future[pg]
+    assert manager._staging_future_to_pg[fut] == pg
+
+    # Wait until PG is ready
+    while not manager.has_resources_ready(resource_request=REQUEST_2_CPU):
+        time.sleep(0.05)
+
+    # PG should now be ready
+    assert manager._request_to_ready_pgs[REQUEST_2_CPU]
+    # PG should not be staged anymore
+    assert not manager._request_to_staged_pgs[REQUEST_2_CPU]
+    # Staging future should not exist anymore
+    assert not manager._pg_to_staging_future
+    assert not manager._staging_future_to_pg
+
+    # Cancel request
+    manager.cancel_resource_request(REQUEST_2_CPU)
+
+    # PG should not be ready anymore
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+    # All PGs should be fully removed
+    assert not manager._pg_to_request
+
+    ####
+    # 2. Request, cancel while staging
+
+    # Stage another PG
+    manager.request_resources(REQUEST_2_CPU)
+    # Cancel request before it's ready
+    manager.cancel_resource_request(REQUEST_2_CPU)
+    # Assert no leftover
+    assert not manager._pg_to_staging_future
+    assert not manager._staging_future_to_pg
+    assert not manager._request_to_staged_pgs[REQUEST_2_CPU]
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+    assert not manager._pg_to_request
+
+    ####
+    # 2. Request, acquire, free
+
+    # Stage another PG
+    manager.request_resources(REQUEST_2_CPU)
+    pg = list(manager._request_to_staged_pgs[REQUEST_2_CPU])[0]
+    # Wait until PG is ready
+    while not manager.has_resources_ready(resource_request=REQUEST_2_CPU):
+        time.sleep(0.05)
+    # Acquire
+    acquired_resources = manager.acquire_resources(resource_request=REQUEST_2_CPU)
+    # Assert no staging/ready leftover
+    assert not manager._pg_to_staging_future
+    assert not manager._staging_future_to_pg
+    assert not manager._request_to_staged_pgs[REQUEST_2_CPU]
+    assert not manager._request_to_ready_pgs[REQUEST_2_CPU]
+    # We still retain this mapping
+    assert manager._pg_to_request
+    # And we keep track of acquired PGs
+    assert pg in manager._acquired_pgs
+
+    # Free PG
+    manager.free_resources(acquired_resources)
+    # State should be cleared now
+    assert not manager._pg_to_request
+    assert not manager._acquired_pgs
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_resource_request.py

@@ -0,0 +1,38 @@
+import pytest
+
+from ray.air.execution.resources.request import ResourceRequest
+
+
+def test_request_same():
+    """Test that resource requests are the same if they share the same properties."""
+
+    assert ResourceRequest([{"CPU": 1}]) == ResourceRequest([{"CPU": 1}])
+
+    # multiple bundles work
+    assert ResourceRequest([{"CPU": 1}, {"CPU": 2}]) == ResourceRequest(
+        [{"CPU": 1}, {"CPU": 2}]
+    )
+
+    # multiple resources work
+    assert ResourceRequest([{"CPU": 1, "GPU": 1}]) == ResourceRequest(
+        [{"CPU": 1, "GPU": 1}]
+    )
+
+    # 0 resources are ignored
+    assert ResourceRequest([{"CPU": 0, "GPU": 1}]) == ResourceRequest([{"GPU": 1}])
+
+    # PACK is implicit
+    assert ResourceRequest([{"CPU": 1}], strategy="PACK") == ResourceRequest(
+        [{"CPU": 1}]
+    )
+
+    # Non match: different strategy
+    assert ResourceRequest([{"CPU": 1}], strategy="PACK") != ResourceRequest(
+        [{"CPU": 1}], strategy="SPREAD"
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_tracked_actor.py

@@ -0,0 +1,366 @@
+import gc
+import threading
+import time
+from collections import Counter
+from typing import Any, Optional, Type
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal import Barrier
+from ray.air.execution._internal.actor_manager import RayActorManager
+
+
+def _raise(exception_type: Type[Exception] = RuntimeError, msg: Optional[str] = None):
+    def _raise_exception(*args, **kwargs):
+        raise exception_type(msg)
+
+    return _raise_exception
+
+
+class Started(RuntimeError):
+    pass
+
+
+class Stopped(RuntimeError):
+    pass
+
+
+class Failed(RuntimeError):
+    pass
+
+
+class Result(RuntimeError):
+    pass
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def cleanup():
+    # Garbage collect at the start
+    # This ensures that all resources are freed up for the upcoming test.
+    gc.collect()
+    yield
+
+
+class Actor:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+    def get_kwargs(self):
+        return self.kwargs
+
+    def task(self, value: Any):
+        return value
+
+
+@ray.remote(num_cpus=4)
+def fn():
+    return True
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("actor_cls", [Actor, ray.remote(Actor)])
+@pytest.mark.parametrize("kill", [False, True])
+def test_start_stop_actor(ray_start_4_cpus, resource_manager_cls, actor_cls, kill):
+    """Test that starting and stopping actors work and invokes a callback.
+
+    - Start an actor
+    - Starting should trigger start callback
+    - Schedule actor task, which should resolve (meaning actor successfully started)
+    - Stop actor, which should resolve and trigger stop callback
+    - Schedule remote fn that takes up all cluster resources. This should resolve,
+      meaning that the actor was stopped successfully.
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    # Start actor, set callbacks
+    tracked_actor = actor_manager.add_actor(
+        cls=actor_cls,
+        kwargs={"key": "val"},
+        resource_request=ResourceRequest([{"CPU": 4}]),
+        on_start=_raise(Started),
+        on_stop=_raise(Stopped),
+        on_error=_raise(Failed),
+    )
+
+    # Actor should be started
+    with pytest.raises(Started):
+        actor_manager.next()
+
+    # Schedule task on actor which should resolve (actor successfully started)
+    actor_manager.schedule_actor_task(
+        tracked_actor, "task", (1,), on_result=_raise(Result)
+    )
+
+    with pytest.raises(Result):
+        actor_manager.next()
+
+    # Now we can assert that there are no CPUS resources available anymore.
+    # Note that actor starting is asynchronous, so we can't assert this right away
+    # - that's why we wait for the actor task to resolve first.
+    assert ray.available_resources().get("CPU", 0.0) == 0, ray.available_resources()
+
+    # Stop actor
+    actor_manager.remove_actor(tracked_actor, kill=kill)
+
+    with pytest.raises(Stopped):
+        actor_manager.next()
+
+    # This task takes up all the cluster resources. It should resolve now that
+    # the actor was terminated.
+    assert ray.get(fn.remote(), timeout=5)
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_start_many_actors(ray_start_4_cpus, resource_manager_cls):
+    """Test that starting more actors than fit onto the cluster works.
+
+    - Request 10 actors
+    - 4 can be started. Assert they are started
+    - Stop 2
+    - Assert 2 are stopped and 2 new ones are started
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    running_actors = []
+    # stats keeps track of started/stopped actors
+    stats = Counter()
+
+    def start_callback(tracked_actor):
+        running_actors.append(tracked_actor)
+        stats["started"] += 1
+
+    def stop_callback(tracked_actor):
+        running_actors.remove(tracked_actor)
+        stats["stopped"] += 1
+
+    # start 10 actors
+    expected_actors = []
+    for i in range(10):
+        tracked_actor = actor_manager.add_actor(
+            cls=Actor,
+            kwargs={"key": "val"},
+            resource_request=ResourceRequest([{"CPU": 1}]),
+            on_start=start_callback,
+            on_stop=stop_callback,
+            on_error=_raise(Failed),
+        )
+        expected_actors.append(tracked_actor)
+
+    # wait for some actor starts
+    for i in range(4):
+        actor_manager.next()
+
+    # we should now have 4 started actors
+    assert stats["started"] == 4
+    assert stats["stopped"] == 0
+    assert len(running_actors) == 4
+    assert set(running_actors) == set(expected_actors[:4])
+
+    # stop 2 actors
+    actor_manager.remove_actor(running_actors[0])
+    actor_manager.remove_actor(running_actors[1])
+
+    # Wait four times, twice for termination, twice for start
+    for i in range(4):
+        actor_manager.next()
+
+    # we should have 4 running actors, 6 started and 2 stopped
+    assert stats["started"] == 6
+    assert stats["stopped"] == 2
+    assert len(running_actors) == 4
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("where", ["init", "fn"])
+def test_actor_fail(ray_start_4_cpus, cleanup, resource_manager_cls, where):
+    """Test that actor failures are handled properly.
+
+    - Start actor that either fails on init or in a task (RayActorError)
+    - Schedule task on actor
+    - Assert that the correct callbacks are called
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    # keep track of failed tasks and actors
+    stats = Counter()
+
+    @ray.remote
+    class FailingActor:
+        def __init__(self, where):
+            self._where = where
+            if self._where == "init":
+                raise RuntimeError("INIT")
+
+        def fn(self):
+            if self._where == "fn":
+                # SystemExit will invoke a RayActorError
+                raise SystemExit
+            return True
+
+    def fail_callback_actor(tracked_actor, exception):
+        stats["failed_actor"] += 1
+
+    def fail_callback_task(tracked_actor, exception):
+        stats["failed_task"] += 1
+
+    # Start actor
+    tracked_actor = actor_manager.add_actor(
+        cls=FailingActor,
+        kwargs={"where": where},
+        resource_request=ResourceRequest([{"CPU": 1}]),
+        on_error=fail_callback_actor,
+    )
+
+    if where != "init":
+        # Wait until it is started. This won't invoke any callback, yet
+        actor_manager.next()
+
+        assert stats["failed_actor"] == 0
+        assert stats["failed_task"] == 0
+
+        # Schedule task
+        actor_manager.schedule_actor_task(
+            tracked_actor, "fn", on_error=fail_callback_task
+        )
+
+    # Yield control and wait for task resolution. This will invoke the callback.
+    actor_manager.next()
+
+    assert stats["failed_actor"] == 1
+    assert stats["failed_task"] == bool(where != "init")
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_stop_actor_before_start(
+    ray_start_4_cpus, tmp_path, cleanup, resource_manager_cls
+):
+    """Test that actor failures are handled properly.
+
+    - Start actor that either fails on init or in a task (RayActorError)
+    - Schedule task on actor
+    - Assert that the correct callbacks are called
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    hang_marker = tmp_path / "hang.txt"
+
+    @ray.remote
+    class HangingActor:
+        def __init__(self):
+            while not hang_marker.exists():
+                time.sleep(0.05)
+
+    tracked_actor = actor_manager.add_actor(
+        HangingActor,
+        kwargs={},
+        resource_request=ResourceRequest([{"CPU": 1}]),
+        on_start=_raise(RuntimeError, "Should not have started"),
+        on_stop=_raise(RuntimeError, "Should not have stopped"),
+    )
+    while not actor_manager.is_actor_started(tracked_actor):
+        actor_manager.next(0.05)
+
+    # Actor started but hasn't triggered on_start, yet
+    actor_manager.remove_actor(tracked_actor)
+    hang_marker.write_text("")
+    while actor_manager.is_actor_started(tracked_actor):
+        actor_manager.next(0.05)
+
+    assert actor_manager.num_live_actors == 0
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+@pytest.mark.parametrize("start_thread", [False, True])
+def test_stop_actor_custom_future(
+    ray_start_4_cpus, tmp_path, cleanup, resource_manager_cls, start_thread
+):
+    """If we pass a custom stop future, the actor should still be shutdown by GC.
+
+    This should also be the case when we start a thread in the background, as we
+    do e.g. in Ray Tune's function runner.
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    hang_marker = tmp_path / "hang.txt"
+
+    actor_name = f"stopping_actor_{resource_manager_cls.__name__}_{start_thread}"
+
+    @ray.remote(name=actor_name)
+    class HangingStopActor:
+        def __init__(self):
+            self._thread = None
+            self._stop_event = threading.Event()
+            if start_thread:
+
+                def entrypoint():
+                    while True:
+                        print("Thread!")
+                        time.sleep(1)
+                        if self._stop_event.is_set():
+                            sys.exit(0)
+
+                self._thread = threading.Thread(target=entrypoint)
+                self._thread.start()
+
+        def stop(self):
+            print("Waiting")
+            while not hang_marker.exists():
+                time.sleep(0.05)
+            self._stop_event.set()
+            print("stopped")
+
+    start_barrier = Barrier(max_results=1)
+    stop_barrier = Barrier(max_results=1)
+
+    tracked_actor = actor_manager.add_actor(
+        HangingStopActor,
+        kwargs={},
+        resource_request=ResourceRequest([{"CPU": 1}]),
+        on_start=start_barrier.arrive,
+        on_stop=stop_barrier.arrive,
+    )
+    while not start_barrier.completed:
+        actor_manager.next(0.05)
+
+    # Actor is alive
+    assert ray.get_actor(actor_name)
+
+    stop_future = actor_manager.schedule_actor_task(tracked_actor, "stop")
+    actor_manager.remove_actor(tracked_actor, kill=False, stop_future=stop_future)
+
+    assert not stop_barrier.completed
+
+    hang_marker.write_text("!")
+
+    while not stop_barrier.completed:
+        actor_manager.next(0.05)
+
+    # Actor should have stopped now and should get cleaned up
+    with pytest.raises(ValueError):
+        ray.get_actor(actor_name)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/execution/test_tracked_actor_task.py

@@ -0,0 +1,122 @@
+from collections import Counter
+
+import pytest
+
+import ray
+from ray.air import ResourceRequest
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.air.execution._internal.actor_manager import RayActorManager
+
+RESOURCE_MANAGERS = [FixedResourceManager, PlacementGroupResourceManager]
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@ray.remote
+class Actor:
+    def foo(self, val, error: bool = False):
+        if error:
+            raise RuntimeError
+        return val
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+def test_resolve(ray_start_4_cpus, resource_manager_cls):
+    """Test that the `on_result` callback is invoked when a task completes.
+
+    - Instantiate global data object
+    - Schedule task that returns a value
+    - The callback writes the returned value to the global data object
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    seen = {"data": 0}
+
+    def result_callback(tracked_actor, result):
+        seen["data"] = result
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.schedule_actor_task(
+        tracked_actor, "foo", (4, False), on_result=result_callback
+    )
+    actor_manager.next()
+    actor_manager.next()
+
+    assert seen["data"] == 4
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+@pytest.mark.parametrize("num_tasks", [1, 10, 100])
+def test_resolve_many(ray_start_4_cpus, resource_manager_cls, num_tasks):
+    """Schedule ``num_tasks`` tasks and wait until ``wait_for_events`` of them resolve.
+
+    Every resolved task will increase a counter by its return value (1).
+    """
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    seen = {"data": 0}
+
+    def result_callback(tracked_actor, result):
+        seen["data"] += result
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.next()
+
+    for i in range(num_tasks):
+        actor_manager.schedule_actor_task(
+            tracked_actor, "foo", (1, False), on_result=result_callback
+        )
+
+    for i in range(num_tasks):
+        actor_manager.next()
+        assert seen["data"] == i + 1
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+def test_error_noop(ray_start_4_cpus, resource_manager_cls):
+    """When no `on_error` callback is specified, errors should be ignored."""
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.schedule_actor_task(tracked_actor, "foo", (1, True))
+    actor_manager.next()
+    actor_manager.next()
+
+
+@pytest.mark.parametrize("resource_manager_cls", RESOURCE_MANAGERS)
+def test_error_custom(ray_start_4_cpus, resource_manager_cls):
+    """When an `on_error` callback is specified, it is invoked."""
+    actor_manager = RayActorManager(resource_manager=resource_manager_cls())
+
+    stats = Counter()
+
+    def error_callback(tracked_actor, exception):
+        stats["exception"] += 1
+
+    tracked_actor = actor_manager.add_actor(
+        cls=Actor, kwargs={}, resource_request=ResourceRequest([{"CPU": 4}])
+    )
+    actor_manager.schedule_actor_task(
+        tracked_actor, "foo", (1, True), on_error=error_callback
+    )
+
+    actor_manager.next()
+    actor_manager.next()
+    assert stats["exception"] == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/mocked_wandb_integration.py

@@ -0,0 +1,136 @@
+from collections import namedtuple
+from dataclasses import dataclass
+from typing import Dict
+from unittest.mock import Mock
+
+from wandb.util import json_dumps_safer
+
+import ray
+from ray.air.integrations.wandb import WandbLoggerCallback, _WandbLoggingActor
+
+
+class Trial(
+    namedtuple(
+        "MockTrial",
+        [
+            "config",
+            "trial_id",
+            "trial_name",
+            "experiment_dir_name",
+            "placement_group_factory",
+            "local_path",
+        ],
+    )
+):
+    def __hash__(self):
+        return hash(self.trial_id)
+
+    def __str__(self):
+        return self.trial_name
+
+
+@dataclass
+class LoggingActorState:
+    args: list
+    kwargs: dict
+    exclude: list
+    logs: list
+    config: dict
+
+
+class _FakeConfig:
+    def __init__(self):
+        self.config = {}
+
+    def update(self, config, *args, **kwargs):
+        self.config.update(config)
+
+
+class _MockWandbAPI:
+    """Thread-safe.
+
+    Note: Not implemented to mock re-init behavior properly. Proceed with caution."""
+
+    def __init__(self):
+        self.logs = []
+        self.config = _FakeConfig()
+
+    def init(self, *args, **kwargs):
+        mock = Mock()
+        mock.args = args
+        mock.kwargs = kwargs
+
+        if "config" in kwargs:
+            self.config.update(kwargs["config"])
+
+        return mock
+
+    def log(self, data, step=None):
+        try:
+            json_dumps_safer(data)
+        except Exception:
+            self.logs.append("serialization error")
+        else:
+            self.logs.append(data)
+
+    def finish(self):
+        pass
+
+    def get_logs(self):
+        return self.logs
+
+    def get_config(self):
+        return self.config.config
+
+
+class _MockWandbLoggingActor(_WandbLoggingActor):
+    _mock_wandb_api_cls = _MockWandbAPI
+
+    def __init__(self, logdir, queue, exclude, to_config, *args, **kwargs):
+        super(_MockWandbLoggingActor, self).__init__(
+            logdir, queue, exclude, to_config, *args, **kwargs
+        )
+        self._wandb = self._mock_wandb_api_cls()
+
+    def get_state(self):
+        return LoggingActorState(
+            args=self.args,
+            kwargs=self.kwargs,
+            exclude=self._exclude,
+            logs=self._wandb.get_logs(),
+            config=self._wandb.get_config(),
+        )
+
+
+class WandbTestExperimentLogger(WandbLoggerCallback):
+    """Wandb logger with mocked Wandb API gateway (one per trial)."""
+
+    _logger_actor_cls = _MockWandbLoggingActor
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._saved_actor_states: Dict["Trial", LoggingActorState] = {}
+
+    def _cleanup_logging_actor(self, trial: "Trial", **kwargs):
+        logging_actor_state: LoggingActorState = ray.get(
+            self._trial_logging_actors[trial].get_state.remote()
+        )
+        self._saved_actor_states[trial] = logging_actor_state
+        super()._cleanup_logging_actor(trial, **kwargs)
+
+    @property
+    def trial_logging_actor_states(self) -> Dict["Trial", LoggingActorState]:
+        return self._saved_actor_states
+
+
+def get_mock_wandb_logger(mock_api_cls=_MockWandbAPI, **kwargs):
+    class MockWandbLoggingActor(_MockWandbLoggingActor):
+        _mock_wandb_api_cls = mock_api_cls
+
+    logger = WandbTestExperimentLogger(
+        project="test_project",
+        api_key="1234",
+        **kwargs,
+    )
+    logger._logger_actor_cls = MockWandbLoggingActor
+    return logger

python/ray/air/tests/test_air_usage.py

@@ -0,0 +1,210 @@
+"""Unit tests for AIR telemetry."""
+
+import json
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pyarrow.fs
+import pytest
+from packaging.version import Version
+
+import ray
+from ray import train, tune
+from ray._common.usage.usage_lib import TagKey
+from ray.air._internal import usage as air_usage
+from ray.air._internal.usage import AirEntrypoint
+from ray.air.integrations import comet, mlflow, wandb
+from ray.train._internal.storage import StorageContext
+from ray.tune.callback import Callback
+from ray.tune.experiment.experiment import Experiment
+from ray.tune.logger import LoggerCallback
+from ray.tune.utils.callback import DEFAULT_CALLBACK_CLASSES
+
+
+def _mock_record_from_module(module, monkeypatch):
+    recorded = {}
+
+    def mock_record_extra_usage_tag(key: TagKey, value: str):
+        recorded[key] = value
+
+    monkeypatch.setattr(
+        module,
+        "record_extra_usage_tag",
+        mock_record_extra_usage_tag,
+    )
+    return recorded
+
+
+@pytest.fixture
+def mock_record(monkeypatch):
+    import ray.air._internal.usage
+
+    yield _mock_record_from_module(ray.air._internal.usage, monkeypatch=monkeypatch)
+
+
+def train_fn(config):
+    train.report({"score": 1})
+
+
+@pytest.fixture
+def tuner(tmp_path):
+    yield tune.Tuner(train_fn, run_config=tune.RunConfig(storage_path=str(tmp_path)))
+
+
+@pytest.fixture
+def trainer(tmp_path):
+    from ray.train.data_parallel_trainer import DataParallelTrainer
+
+    yield DataParallelTrainer(
+        train_loop_per_worker=train_fn,
+        scaling_config=train.ScalingConfig(num_workers=2),
+        run_config=train.RunConfig(storage_path=str(tmp_path)),
+    )
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "storage_path_filesystem_expected",
+    [
+        ("/tmp/test", None, "local"),
+        ("s3://", None, "s3"),
+        ("gs://test", None, "gcs"),
+        ("mock://test", None, "mock"),
+        ("test", pyarrow.fs.LocalFileSystem(), "custom"),
+    ],
+)
+def test_tag_storage_type(storage_path_filesystem_expected, mock_record, monkeypatch):
+    # Don't write anything to storage for the test.
+    monkeypatch.setattr(StorageContext, "_create_validation_file", lambda _: None)
+    monkeypatch.setattr(StorageContext, "_check_validation_file", lambda _: None)
+
+    storage_path, storage_filesystem, expected = storage_path_filesystem_expected
+
+    if Version(pyarrow.__version__) < Version("17.0.0") and storage_path.startswith(
+        "gs://"
+    ):
+        pytest.skip("GCS support requires pyarrow >= 17.0.0")
+
+    storage = StorageContext(
+        storage_path=storage_path,
+        experiment_dir_name="test",
+        storage_filesystem=storage_filesystem,
+    )
+    air_usage.tag_storage_type(storage)
+    assert mock_record[TagKey.AIR_STORAGE_CONFIGURATION] == expected
+
+
+class _CustomLoggerCallback(LoggerCallback):
+    pass
+
+
+class _CustomCallback(Callback):
+    pass
+
+
+_TEST_CALLBACKS = [
+    wandb.WandbLoggerCallback,
+    mlflow.MLflowLoggerCallback,
+    comet.CometLoggerCallback,
+    _CustomLoggerCallback,
+    _CustomLoggerCallback,
+    _CustomCallback,
+]
+
+
+def test_tag_setup_wandb(mock_record):
+    from ray.air.integrations.wandb import _setup_wandb
+
+    with patch.dict(os.environ, {wandb.WANDB_MODE_ENV_VAR: "disabled"}):
+        _setup_wandb(trial_id="a", trial_name="b", config={}, _wandb=MagicMock())
+    assert mock_record[TagKey.AIR_SETUP_WANDB_INTEGRATION_USED] == "1"
+
+
+def test_tag_setup_mlflow(mock_record, monkeypatch):
+    from ray.air.integrations.mlflow import setup_mlflow
+
+    monkeypatch.setattr(ray.air.integrations.mlflow, "_MLflowLoggerUtil", MagicMock())
+    setup_mlflow()
+    assert mock_record[TagKey.AIR_SETUP_MLFLOW_INTEGRATION_USED] == "1"
+
+
+@pytest.mark.parametrize(
+    "callback_classes_expected",
+    [
+        (None, None),
+        ([], None),
+        ([lambda: None], None),
+        (
+            DEFAULT_CALLBACK_CLASSES,
+            {cls.__name__: 1 for cls in DEFAULT_CALLBACK_CLASSES},
+        ),
+        (
+            _TEST_CALLBACKS,
+            {
+                "WandbLoggerCallback": 1,
+                "MLflowLoggerCallback": 1,
+                "CometLoggerCallback": 1,
+                "CustomLoggerCallback": 2,
+                "CustomCallback": 1,
+            },
+        ),
+    ],
+)
+def test_tag_callbacks(mock_record, callback_classes_expected):
+    callback_classes, expected = callback_classes_expected
+
+    callbacks = (
+        [callback_cls() for callback_cls in callback_classes]
+        if callback_classes
+        else None
+    )
+
+    air_usage.tag_callbacks(callbacks)
+
+    callback_usage_str = mock_record.pop(TagKey.AIR_CALLBACKS, None)
+    callback_counts = json.loads(callback_usage_str) if callback_usage_str else None
+    assert callback_counts == expected
+
+
+def test_tag_env_vars(ray_start_4_cpus, mock_record, tuner):
+    """Test that env vars are recorded properly, and arbitrary user environment
+    variables are ignored."""
+    env_vars_to_record = {
+        "TUNE_GLOBAL_CHECKPOINT_S": "20",
+        "TUNE_MAX_PENDING_TRIALS_PG": "1",
+    }
+    untracked_env_vars = {"RANDOM_USER_ENV_VAR": "asdf"}
+
+    with patch.dict(os.environ, {**env_vars_to_record, **untracked_env_vars}):
+        tuner.fit()
+
+    recorded_env_vars = json.loads(mock_record[TagKey.AIR_ENV_VARS])
+    assert sorted(env_vars_to_record) == sorted(recorded_env_vars)
+
+
+@pytest.mark.parametrize("entrypoint", list(AirEntrypoint))
+def test_tag_air_entrypoint(ray_start_4_cpus, mock_record, entrypoint, tuner, trainer):
+    if entrypoint == AirEntrypoint.TUNE_RUN:
+        tune.run(train_fn)
+    elif entrypoint == AirEntrypoint.TUNE_RUN_EXPERIMENTS:
+        experiment_spec = Experiment("experiment", train_fn)
+        tune.run_experiments(experiments=experiment_spec)
+    elif entrypoint == AirEntrypoint.TUNER:
+        tuner.fit()
+    elif entrypoint == AirEntrypoint.TRAINER:
+        trainer.fit()
+
+    assert mock_record[TagKey.AIR_ENTRYPOINT] == entrypoint.value
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/air/tests/test_errors.py

@@ -0,0 +1,284 @@
+"""
+This test suite covers error handling and propagation in Ray Train/Tune.
+
+There are two main error types to test:
+1. Trainable errors: These happen in the remote actor itself.
+    -> Within this, we should test:
+        - fail_fast=True/False/'raise'
+        - AIR Trainer w/o Tuner, AIR Trainer w/ Tuner, Tuner w/ function trainable
+2. Tune driver errors: These happen in the Tune event-handling loop.
+    -> Within this, we should test:
+        - Errors occurring at different points in the Tune loop
+          (on_trial_result, on_checkpoint, on_step_begin, etc.)
+
+These tests should:
+- Assert how errors from the trainable/Trainer get propagated to the user.
+- Assert how errors from the Tune driver get propagated to the user.
+"""
+
+import gc
+import threading
+import time
+from tempfile import TemporaryDirectory
+
+import pytest
+
+import ray
+from ray import train, tune
+from ray._common.test_utils import wait_for_condition
+from ray._raylet import GcsClient
+from ray.cluster_utils import Cluster
+from ray.core.generated import autoscaler_pb2
+from ray.tests.conftest import *  # noqa
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.trainer import BaseTrainer, TrainingFailedError
+from ray.tune import TuneError, Tuner
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4, configure_logging=False)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def gc_collect():
+    # Make sure to cleanup as much as possible between
+    # unit tests that share a Ray session
+    yield
+    gc.collect()
+
+
+@pytest.fixture
+def cluster_setup(ray_start_cluster_head: Cluster):
+    # Sets up a cluster with 3 nodes: head node + 2 workers
+    cluster = ray_start_cluster_head
+    nodes = []
+    nodes.append(cluster.add_node(resources={"worker1": 1, "cpu": 1, "coordinator": 1}))
+    nodes.append(cluster.add_node(resources={"worker2": 1, "cpu": 1}))
+    cluster.wait_for_nodes()
+
+    @ray.remote
+    def get_node_id():
+        return ray.get_runtime_context().get_node_id()
+
+    worker1_node_id = ray.get(get_node_id.options(resources={"worker1": 1}).remote())
+    worker2_node_id = ray.get(get_node_id.options(resources={"worker2": 1}).remote())
+    wait_for_condition(
+        lambda: len({node["NodeID"] for node in ray.nodes() if (node["Alive"])}) == 3
+    )
+
+    yield cluster, nodes, [
+        worker1_node_id,
+        worker2_node_id,
+    ]
+
+
+class _TestSpecificError(RuntimeError):
+    pass
+
+
+class FailingCallback(tune.Callback):
+    def __init__(self, error_on: str):
+        self.error_on = error_on
+
+    def on_trial_result(self, *args, **kwargs):
+        if self.error_on == "on_trial_result":
+            raise _TestSpecificError(f"Failing on {self.error_on}!")
+
+
+class FailingTrainer(BaseTrainer):
+    def training_loop(self) -> None:
+        raise _TestSpecificError("There is an error in trainer!")
+
+
+def passing_fn(config):
+    # Trigger all the driver events (on_checkpoint, on_trial_save, etc.)
+    with TemporaryDirectory() as tmpdir:
+        train.report({"score": 1}, checkpoint=train.Checkpoint.from_directory(tmpdir))
+
+
+def failing_fn(config):
+    raise _TestSpecificError("Failing!")
+
+
+trainable_map = {
+    "function": failing_fn,
+    "trainer": FailingTrainer(),
+}
+
+
+@pytest.mark.parametrize("fail_fast", [False, True, "raise"])
+def test_trainable_error_with_tuner(ray_start_4_cpus, fail_fast):
+    tuner = Tuner(
+        trainable=failing_fn,
+        run_config=tune.RunConfig(
+            name=f"tuner_errors-fail_fast={fail_fast}",
+            failure_config=tune.FailureConfig(fail_fast=fail_fast),
+        ),
+        tune_config=tune.TuneConfig(num_samples=2),
+    )
+
+    if fail_fast is False:
+        # Both trials should complete with an error.
+        results = tuner.fit()
+        assert len(results) == 2
+        for i in range(2):
+            assert results[i].error
+    elif fail_fast is True:
+        # The first trial errors -> the experiment finishes immediately.
+        results = tuner.fit()
+        errors = [result.error for result in results if result.error]
+        assert len(errors) == 1
+    elif fail_fast == "raise":
+        # The original error gets raised to the user
+        with pytest.raises(_TestSpecificError):
+            tuner.fit()
+
+
+@pytest.mark.parametrize("fail_fast", [False, True, "raise"])
+def test_trainable_error_with_trainer(ray_start_4_cpus, tmp_path, fail_fast):
+    name = f"test_trainer_errors-fail_fast={fail_fast}"
+    trainer = FailingTrainer(
+        run_config=train.RunConfig(
+            storage_path=str(tmp_path),
+            name=name,
+            failure_config=train.FailureConfig(fail_fast=fail_fast),
+        ),
+        scaling_config=train.ScalingConfig(num_workers=1),
+    )
+
+    if fail_fast in [False, True]:
+        # There is only 1 "trial" for a Trainer,
+        # so fail_fast = True/False doesn't change the behavior
+        # In both cases, the error should get wrapped and raised.
+        with pytest.raises(TrainingFailedError) as exc_info:
+            trainer.fit()
+
+        # The cause of the error should be the trainable error
+        assert isinstance(exc_info.value.__cause__, _TestSpecificError)
+
+        assert TrainingFailedError._RESTORE_MSG.format(
+            trainer_cls_name="FailingTrainer", path=str(tmp_path / name)
+        ) in str(exc_info.value)
+        assert TrainingFailedError._FAILURE_CONFIG_MSG in str(exc_info.value)
+
+    elif fail_fast == "raise":
+        # The original error gets raised to the user
+        with pytest.raises(_TestSpecificError):
+            trainer.fit()
+
+
+# TODO(ml-team): Test all the driver hooks once driver error propagation is fixed
+
+
+@pytest.mark.parametrize("error_on", ["on_trial_result"])
+def test_driver_error_with_tuner(ray_start_4_cpus, error_on):
+    tuner = Tuner(
+        trainable=passing_fn,
+        run_config=tune.RunConfig(
+            name=f"test_driver_errors_with_tuner-error_on={error_on}",
+            callbacks=[FailingCallback(error_on=error_on)],
+        ),
+    )
+
+    # All driver errors should get propagated to the user in the same way
+    with pytest.raises(TuneError) as exc_info:
+        tuner.fit()
+
+    # TODO(ml-team): Assert the cause error type once driver error propagation is fixed
+    assert "_TestSpecificError" in str(exc_info.value)
+
+
+@pytest.mark.parametrize("error_at_level", ["worker", "coordinator"])
+def test_preemption_handling(
+    cluster_setup,
+    tmp_path,
+    error_at_level: str,
+):
+    """Integration test for node preemption handling in Ray Train/Tune.
+    Even though `max_failures=0`, preemption errors should still be retried."""
+    cluster, nodes, node_ids = cluster_setup
+    # node 1 = coordinator and worker, node 2 = worker
+    coordinator_node, worker_node = nodes
+    coordinator_node_id, worker_node_id = node_ids
+
+    num_workers = 2
+    tmp_path.joinpath("markers").mkdir()
+
+    def train_fn(config):
+        checkpoint = train.get_checkpoint()
+        start_iter = 0
+        if checkpoint:
+            start_iter = load_dict_checkpoint(checkpoint)["iter"] + 1
+            print(f"Restored at iter = {start_iter}")
+
+        for iter in range(start_iter, 6):
+            with create_dict_checkpoint({"iter": iter}) as checkpoint:
+                ray.train.report({"iter": iter}, checkpoint=checkpoint)
+
+            if iter == 2:
+                # Write a "done marker" to tell the driver to simulate a preemption.
+                tmp_path.joinpath(
+                    "markers", str(ray.train.get_context().get_world_rank())
+                ).touch()
+                # Await execution.
+                time.sleep(120)
+
+    def launch_training():
+        trainer = DataParallelTrainer(
+            train_loop_per_worker=train_fn,
+            scaling_config=train.ScalingConfig(
+                num_workers=num_workers,
+                trainer_resources={"coordinator": 1},
+                resources_per_worker={"cpu": 1},  # worker2 and worker3
+            ),
+            run_config=train.RunConfig(
+                storage_path=str(tmp_path),
+                name="test_preemption_error",
+                failure_config=train.FailureConfig(fail_fast=False, max_failures=0),
+            ),
+        )
+        result = trainer.fit()
+        assert result.metrics["iter"] == 5
+
+    t = threading.Thread(target=launch_training)
+    t.start()
+
+    # Wait until the workers are ready for preemption (after a few checkpoints).
+    while len(list(tmp_path.joinpath("markers").glob("*"))) < num_workers:
+        time.sleep(0.5)
+
+    if error_at_level == "coordinator":
+        node, node_id = coordinator_node, coordinator_node_id
+    elif error_at_level == "worker":
+        node, node_id = worker_node, worker_node_id
+    else:
+        raise NotImplementedError(f"Invalid error_at_level = {error_at_level}")
+
+    # Preempt a node.
+    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
+    print("Draining node...")
+    is_accepted, _ = gcs_client.drain_node(
+        node_id,
+        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
+        "preemption",
+        0,
+    )
+    assert is_accepted
+    print("Killing node...")
+    cluster.remove_node(node, allow_graceful=True)
+    print("Adding new node..")  # so that the job can be rescheduled
+    # New node can replace a preempted coordinator or worker
+    # NOTE: `cluster.add_node` only works in the main thread.
+    cluster.add_node(resources={"coordinator": 1, "cpu": 1})
+    t.join()  # Assert no errors during training.
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))

python/ray/air/tests/test_experiment_restore.py

@@ -0,0 +1,232 @@
+import json
+import os
+import shutil
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from ray.tune.analysis import ExperimentAnalysis
+from ray.tune.result_grid import ResultGrid
+
+_RUN_SCRIPT_FILENAME = "_test_experiment_restore_run.py"
+
+
+def _kill_process_if_needed(
+    process: subprocess.Popen, timeout_s: float = 10, poll_interval_s: float = 1.0
+):
+    """Kills a process if it hasn't finished in `timeout_s` seconds.
+    Polls every `poll_interval_s` seconds to check if the process is still running."""
+    kill_timeout = time.monotonic() + timeout_s
+    while process.poll() is None and time.monotonic() < kill_timeout:
+        time.sleep(poll_interval_s)
+    if process.poll() is None:
+        process.terminate()
+
+
+def _print_message(message):
+    sep = "=" * 50
+    print(f"\n{sep}\n{message}\n{sep}\n")
+
+
+@pytest.mark.parametrize("runner_type", ["tuner", "trainer"])
+def test_experiment_restore(tmp_path, runner_type):
+    """
+    This is an integration stress test for experiment restoration.
+
+
+    Test setup:
+
+    - For Tuner.restore:
+        - 8 trials, with a max of 2 running concurrently (--> 4 rounds of trials)
+        - Each iteration takes 0.5 seconds
+        - Each trial runs for 8 iterations --> 4 seconds
+        - Each round of 2 trials should take 4 seconds
+        - Without any interrupts/restoration:
+            - Minimum runtime: 4 rounds * 4 seconds / round = 16 seconds
+        - The test will stop the script with a SIGINT at a random time between
+        6-10 iterations each restore.
+
+    - For Trainer.restore:
+        - 1 trial with 4 workers
+        - Each iteration takes 0.5 seconds
+        - Runs for 32 iterations --> Minimum runtime = 16 seconds
+        - The test will stop the script with a SIGINT at a random time between
+        6-10 iterations after each restore.
+
+    Requirements:
+    - Req 1: Training progress persisted
+        - The experiment should progress monotonically.
+        (The training iteration shouldn't go backward at any point)
+        - Trials shouldn't start from scratch.
+    - Req 2: Searcher state saved/restored correctly
+    - Req 3: Callback state saved/restored correctly
+    """
+
+    np.random.seed(2023)
+
+    script_path = Path(__file__).parent / _RUN_SCRIPT_FILENAME
+
+    # Args to pass into the script as environment variables
+    exp_name = f"{runner_type}_restore_integration_test"
+    callback_dump_file = tmp_path / f"{runner_type}-callback_dump_file.json"
+    storage_path = tmp_path / "ray_results"
+    if storage_path.exists():
+        shutil.rmtree(storage_path)
+
+    csv_file = str(tmp_path / "dummy_data.csv")
+    dummy_df = pd.DataFrame({"x": np.arange(128), "y": 2 * np.arange(128)})
+    dummy_df.to_csv(csv_file)
+
+    run_started_marker = tmp_path / "run_started_marker"
+
+    time_per_iter_s = 0.5
+    max_concurrent = 2
+
+    if runner_type == "tuner":
+        iters_per_trial = 8
+        num_trials = 8
+    elif runner_type == "trainer":
+        iters_per_trial = 32
+        num_trials = 1
+
+    total_iters = iters_per_trial * num_trials
+
+    env = os.environ.copy()
+    env.update(
+        {
+            "RUNNER_TYPE": runner_type,
+            "STORAGE_PATH": str(storage_path),
+            "EXP_NAME": exp_name,
+            "CALLBACK_DUMP_FILE": str(callback_dump_file),
+            "RUN_STARTED_MARKER": str(run_started_marker),
+            "TIME_PER_ITER_S": str(time_per_iter_s),
+            "ITERATIONS_PER_TRIAL": str(iters_per_trial),
+            "NUM_TRIALS": str(num_trials),
+            "MAX_CONCURRENT_TRIALS": str(max_concurrent),
+            "CSV_DATA_FILE": csv_file,
+        }
+    )
+
+    # Variables used in the loop
+    return_code = None
+    total_runtime = 0
+    run_iter = 0
+    progress = 0
+    progress_history = []
+
+    poll_interval_s = 0.1
+    test_start_time = time.monotonic()
+
+    while True:
+        run_started_marker.write_text("", encoding="utf-8")
+
+        run = subprocess.Popen([sys.executable, script_path], env=env)
+        run_iter += 1
+
+        _print_message(f"Started run #{run_iter} w/ PID = {run.pid}")
+
+        # Start the timer after the first trial has entered its training loop.
+        while run.poll() is None and run_started_marker.exists():
+            time.sleep(poll_interval_s)
+
+        # If the run already finished, then exit immediately.
+        if run.poll() is not None:
+            return_code = run.poll()
+            break
+
+        timeout_s = np.random.uniform(6 * time_per_iter_s, 10 * time_per_iter_s)
+
+        _print_message(
+            "Training has started...\n"
+            f"Interrupting after {timeout_s:.2f} seconds\n"
+            f"Currently at {total_runtime:.2f} seconds"
+        )
+
+        # Sleep for a random amount of time, then stop the run.
+        start_time = time.monotonic()
+        time.sleep(timeout_s)
+        total_runtime += time.monotonic() - start_time
+
+        return_code = run.poll()
+        if return_code is None:
+            # Send "SIGINT" to stop the run
+            _print_message(f"Sending SIGUSR1 to run #{run_iter} w/ PID = {run.pid}")
+            run.send_signal(signal.SIGUSR1)
+
+            # Make sure the process is stopped forcefully after a timeout.
+            _kill_process_if_needed(run)
+        else:
+            _print_message("Run has already terminated!")
+            break
+
+        # Check up on the results.
+        results = ResultGrid(ExperimentAnalysis(str(storage_path / exp_name)))
+        iters = [result.metrics.get("training_iteration", 0) for result in results]
+        progress = sum(iters) / total_iters
+        progress_history.append(progress)
+        _print_message(
+            f"Number of trials = {len(results)}\n"
+            f"% completion = {progress} ({sum(iters)} iters / {total_iters})\n"
+            f"Currently at {total_runtime:.2f} seconds"
+        )
+
+    _print_message(
+        f"Total number of restorations = {run_iter}\n"
+        f"Total runtime = {total_runtime:.2f}\n"
+        f"Return code = {return_code}"
+    )
+    test_end_time = time.monotonic()
+
+    assert progress == 1.0
+
+    # The script shouldn't have errored. (It should have finished by this point.)
+    assert return_code == 0, (
+        f"The script errored with return code: {return_code}.\n"
+        f"Check the `{_RUN_SCRIPT_FILENAME}` script for any issues. "
+    )
+
+    # Req 1: training progress persisted
+    # Check that progress increases monotonically (we never go backwards/start from 0)
+    assert np.all(np.diff(progress_history) >= 0), (
+        "Expected progress to increase monotonically. Instead, got:\n"
+        "{progress_history}"
+    )
+
+    # Req 2: searcher state
+    results = ResultGrid(ExperimentAnalysis(str(storage_path / exp_name)))
+    # Check that all trials have unique ids assigned by the searcher (if applicable)
+    ids = [result.config.get("id", -1) for result in results]
+    ids = [id for id in ids if id >= 0]
+    if ids:
+        assert sorted(ids) == list(range(1, num_trials + 1)), (
+            "Expected the searcher to assign increasing id for each trial, but got:"
+            f"{ids}"
+        )
+
+    # Req 3: callback state
+    with open(callback_dump_file, "r") as f:
+        callback_state = json.load(f)
+
+    trial_iters = callback_state["trial_iters"]
+    for iters in trial_iters.values():
+        # Check that the callback has data for each trial, for all iters
+        # NOTE: There may be some duplicate data, due to the fact that
+        # the callback will be updated on every `on_trial_result` hook,
+        # but the trial may crash before the corresponding checkpoint gets processed.
+        assert sorted(set(iters)) == list(
+            range(1, iters_per_trial + 1)
+        ), f"Expected data from all iterations, but got: {iters}"
+
+    _print_message(f"Success! Test took {test_end_time - test_start_time:.2f} seconds.")
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_integration_comet.py

@@ -0,0 +1,366 @@
+import unittest
+from collections import namedtuple
+from unittest.mock import patch
+
+from ray.air.integrations.comet import CometLoggerCallback
+
+
+class MockTrial(
+    namedtuple("MockTrial", ["config", "trial_name", "trial_id", "logdir"])
+):
+    def __hash__(self):
+        return hash(self.trial_id)
+
+    def __str__(self):
+        return self.trial_name
+
+
+class InitializationTests(unittest.TestCase):
+    def setUp(self):
+        self.logger = CometLoggerCallback()
+
+    def test_class_variable_to_instance(self):
+        """Test that class variables get properly assigned to instance
+        variables.
+        """
+        logger = self.logger
+        self.assertEqual(logger._to_exclude, logger._exclude_results)
+        self.assertEqual(logger._to_system, logger._system_results)
+        self.assertEqual(logger._to_other, logger._other_results)
+        self.assertEqual(logger._to_episodes, logger._episode_results)
+
+    def test_configure_experiment_defaults(self):
+        """Test CometLoggerCallback._configure_experiment_defaults."""
+        logger = self.logger
+
+        # Test that autologging features are properly disabled
+        exclude = CometLoggerCallback._exclude_autolog
+        for option in exclude:
+            self.assertFalse(logger.experiment_kwargs.get(option))
+        del logger
+
+        # Don't disable logging if user overwrites defaults by passing in args
+        for include_option in exclude:
+            # This unpacks to become e.g. CometLoggerCallback(log_env_cpu=True)
+            logger = CometLoggerCallback(**{include_option: True})
+            for option in exclude:
+                if option == include_option:
+                    self.assertTrue(logger.experiment_kwargs.get(option))
+                else:
+                    self.assertFalse(logger.experiment_kwargs.get(option))
+
+
+class HelperMethodTests(unittest.TestCase):
+    def setUp(self):
+        self.logger = CometLoggerCallback()
+
+    def test_check_key_name(self):
+
+        logger = self.logger
+        # Return True when key == item
+        self.assertTrue(logger._check_key_name("name", "name"))
+        # Return True when key.startswith(item + "/")
+        self.assertTrue(logger._check_key_name("name/", "name"))
+        # Return False when item.startswith(key + "/")
+        self.assertFalse(logger._check_key_name("name", "name/"))
+        # Return False when key != item and not key.startswith(item."/")
+        self.assertFalse(logger._check_key_name("name", "x"))
+
+
+@patch("comet_ml.OfflineExperiment")
+@patch("comet_ml.Experiment")
+class OnlineVsOfflineTests(unittest.TestCase):
+    def setUp(self):
+        self.loggers = {
+            "online": CometLoggerCallback(),
+            "offline": CometLoggerCallback(online=False),
+        }
+
+        self.trial = MockTrial({"p1": 1}, "trial_1", 1, "artifact")
+
+    def test_online_dispatch(self, experiment, offline_experiment):
+
+        # To start, there should be no experiments
+        experiment.assert_not_called()
+        offline_experiment.assert_not_called()
+
+        # Start online experiment
+        logger = self.loggers["online"]
+        logger.log_trial_start(self.trial)
+
+        # Check that Experiment was called and OfflineExperiment was not
+        experiment.assert_called_once()
+        offline_experiment.assert_not_called()
+
+    def test_offline_dispatch(self, experiment, offline_experiment):
+
+        # To start, there should be no experiments
+        experiment.assert_not_called()
+        offline_experiment.assert_not_called()
+
+        # Start online experiment
+        logger = self.loggers["offline"]
+        logger.log_trial_start(self.trial)
+
+        # Check that Experiment was called and OfflineExperiment was not
+        experiment.assert_not_called()
+        offline_experiment.assert_called_once()
+
+
+@patch("comet_ml.OfflineExperiment")
+@patch("comet_ml.Experiment")
+class LogTrialStartTest(unittest.TestCase):
+    def setUp(self):
+        self.loggers = {
+            "online": CometLoggerCallback(),
+            "offline": CometLoggerCallback(online=False),
+        }
+
+        self.trials = [
+            MockTrial({"p1": 1}, "trial_1", 1, "artifact"),
+            MockTrial({"p1": 2}, "trial_2", 1, "artifact"),
+        ]
+
+    def test_existing_trialexperiment(self, experiment, offline_experiment):
+
+        mocks = {"online": experiment, "offline": offline_experiment}
+        for option in ["online", "offline"]:
+            logger = self.loggers[option]
+            mock = mocks[option]
+
+            # This should create an experiment
+            logger.log_trial_start(self.trials[0])
+            mock.assert_called_once()
+
+            # This should NOT create an experiment because it's the same trial
+            logger.log_trial_start(self.trials[0])
+            mock.assert_called_once()
+
+            # This should create another new experiment
+            logger.log_trial_start(self.trials[1])
+
+            # Number of times the mock was called
+            num_calls = len(mock.call_args_list)
+
+            # Assert that Experiment/OfflineExperiment was called twice
+            self.assertEqual(num_calls, 2)
+
+    def test_set_global_experiment(self, experiment, offline_experiment):
+        for option in ["online", "offline"]:
+            logger = self.loggers[option]
+            with patch("comet_ml.config.set_global_experiment") as mock:
+                logger.log_trial_start(self.trials[0])
+                mock.assert_called_with(None)
+                mock.assert_called_once()
+                mock.reset_mock()
+
+    def test_experiment_addtags(self, experiment, offline_experiment):
+        logger = self.loggers["online"]
+        logger.log_trial_start(self.trials[0])
+        experiment.return_value.add_tags.assert_called_with(logger.tags)
+
+    def test_experiment_setname(self, experiment, offline_experiment):
+        logger = self.loggers["online"]
+        trial = self.trials[0]
+        logger.log_trial_start(trial)
+        experiment.return_value.set_name.assert_called_with(trial.trial_name)
+
+    def test_experiment_logparams(self, experiment, offline_experiment):
+        logger = self.loggers["online"]
+        trial = self.trials[0]
+        logger.log_trial_start(trial)
+        config = trial.config.copy()
+        config.pop("callbacks", None)
+        experiment.return_value.log_parameters.assert_called_with(config)
+
+
+class ExperimentKwargsTest(unittest.TestCase):
+    @patch("comet_ml.Experiment")
+    def test_kwargs_passthrough(self, experiment):
+        """Test that additional keyword arguments to CometLoggerCallback get
+        passed through to comet_ml.Experiment on log_trial_start
+        """
+        experiment_kwargs = {"kwarg_1": "val_1"}
+        logger = CometLoggerCallback(**experiment_kwargs)
+        trial = MockTrial({"parameter": 1}, "trial2", 1, "artifact")
+        logger.log_trial_start(trial)
+
+        # These are the default kwargs that get passed to create the experiment
+        expected_kwargs = {kwarg: False for kwarg in logger._exclude_autolog}
+        expected_kwargs.update(experiment_kwargs)
+
+        experiment.assert_called_with(**expected_kwargs)
+
+
+@patch("comet_ml.Experiment")
+class LogTrialResultTests(unittest.TestCase):
+    """
+    * test log_others logs
+    * test log_system logs
+    * test log_curve logs
+    """
+
+    def setUp(self):
+        self.logger = CometLoggerCallback()
+        self.trials = [
+            MockTrial({"p1": 1}, "trial_1", 1, "artifact"),
+            MockTrial({"p1": 2}, "trial_2", 1, "artifact"),
+        ]
+        self.result = {
+            "config": {"p1": 1},
+            "node_ip": "0.0.0.0",
+            "hostname": "hostname_val",
+            "pid": "1234",
+            "date": "2000-01-01",
+            "experiment_id": "1234",
+            "trial_id": 1,
+            "experiment_tag": "tag1",
+            "hist_stats/episode_reward": [1, 0, 1, -1, 0, 1],
+            "hist_stats/episode_lengths": [1, 2, 3, 4, 5, 6],
+            "metric1": 0.8,
+            "metric2": 1,
+            "metric3": None,
+            "training_iteration": 0,
+        }
+
+    def test_log_parameters(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+        result = self.result.copy()
+
+        # Check parameters are logged properly.
+        logger.log_trial_result(1, trial, self.result)
+
+        config_update = result.copy().pop("config", {})
+        config_update.pop("callbacks", None)  # Remove callbacks
+
+        experiment.return_value.log_parameters.assert_any_call(config_update)
+
+    def test_log_metrics(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+        result = self.result.copy()
+        step = result["training_iteration"]
+
+        logger.log_trial_result(1, trial, self.result)
+        result_metrics = {
+            "metric1": 0.8,
+            "metric2": 1,
+            "metric3": None,
+            "training_iteration": 0,
+        }
+
+        method = experiment.return_value.log_metrics
+        method.assert_any_call(result_metrics, step=step)
+
+    def test_log_other(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+        result = self.result.copy()
+
+        logger.log_trial_result(1, trial, result)
+        result_other = {
+            "experiment_id": "1234",
+            "trial_id": 1,
+            "experiment_tag": "tag1",
+        }
+        method = experiment.return_value.log_others
+
+        # for k,v in result_other.items():
+        method.assert_any_call(result_other)
+
+    def test_log_system(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+        result = self.result.copy()
+
+        logger.log_trial_result(1, trial, result)
+        result_system = {
+            "node_ip": "0.0.0.0",
+            "hostname": "hostname_val",
+            "pid": "1234",
+            "date": "2000-01-01",
+        }
+        method = experiment.return_value.log_system_info
+        for k, v in result_system.items():
+            method.assert_any_call(k, v)
+
+    def test_log_curve(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+
+        # Check parameters are logged properly.
+        result = self.result
+        step = result["training_iteration"]
+        logger.log_trial_result(1, trial, result)
+
+        results_curve = {
+            "hist_stats/episode_reward": [1, 0, 1, -1, 0, 1],
+            "hist_stats/episode_lengths": [1, 2, 3, 4, 5, 6],
+        }
+
+        method = experiment.return_value.log_curve
+        print(method.call_args_list)
+        for k, v in results_curve.items():
+
+            method.assert_any_call(k, x=range(len(v)), y=v, step=step)
+
+
+@patch("comet_ml.Experiment")
+class LogTrialEndTests(unittest.TestCase):
+    def setUp(self):
+        self.logger = CometLoggerCallback()
+        self.trials = [
+            MockTrial({"p1": 1}, "trial_1", 1, "artifact"),
+            MockTrial({"p1": 2}, "trial_2", 2, "artifact"),
+            MockTrial({"p1": 2}, "trial_3", 3, "artifact"),
+        ]
+
+    def test_not_started_exception(self, experiment):
+        logger = self.logger
+        with self.assertRaises(KeyError):
+            logger.log_trial_end(self.trials[0])
+
+    def test_repeat_throws_error(self, experiment):
+        logger = self.logger
+        trial = self.trials[0]
+
+        logger.log_trial_start(trial)
+        logger.log_trial_end(trial)
+        with self.assertRaises(KeyError):
+            logger.log_trial_end(trial)
+
+    def test_log_trial_end(self, experiment):
+        logger = self.logger
+        trials = self.trials
+        method = experiment.return_value.end
+
+        # Should not have ended yet
+        method.assert_not_called()
+
+        for trial in trials:
+            logger.log_trial_start(trial)
+            logger.log_trial_end(trial)
+
+        self.assertEqual(len(method.call_args_list), len(trials))
+
+    def test_del(self, experiment):
+        logger = self.logger
+
+        for trial in self.trials:
+            logger.log_trial_start(trial)
+
+        end = experiment.return_value.end
+        end.assert_not_called()
+
+        logger.__del__()
+
+        self.assertEqual(len(end.call_args_list), len(self.trials))
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_integration_mlflow.py

@@ -0,0 +1,442 @@
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+from collections import namedtuple
+from unittest.mock import MagicMock, patch
+
+import pytest
+from mlflow.tracking import MlflowClient
+
+import ray
+from ray._private.dict import flatten_dict
+from ray.air._internal.mlflow import _MLflowLoggerUtil
+from ray.air.integrations.mlflow import MLflowLoggerCallback, _NoopModule, setup_mlflow
+from ray.train.torch import TorchTrainer
+from ray.tune import Tuner
+
+
+class MockTrial(
+    namedtuple("MockTrial", ["config", "trial_name", "trial_id", "local_path"])
+):
+    def __hash__(self):
+        return hash(self.trial_id)
+
+    def __str__(self):
+        return self.trial_name
+
+
+class Mock_MLflowLoggerUtil(_MLflowLoggerUtil):
+    def save_artifacts(self, dir, run_id):
+        self.artifact_saved = True
+        self.artifact_info = {"dir": dir, "run_id": run_id}
+
+
+def clear_env_vars():
+    os.environ.pop("MLFLOW_EXPERIMENT_NAME", None)
+    os.environ.pop("MLFLOW_EXPERIMENT_ID", None)
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    """Automatically start and stop Ray for each test."""
+    ray.init(num_cpus=4)
+    yield
+    ray.shutdown()
+
+
+def test_setup_mlflow_in_train_worker(ray_start_4_cpus):
+    """Test that setup_mlflow works in a Train worker."""
+
+    def train_func(config):
+        setup_mlflow(
+            experiment_name="test_exp",
+            create_experiment_if_not_exists=True,
+        )
+
+    trainer = TorchTrainer(train_func)
+    trainer.fit()
+
+
+def test_setup_mlflow_in_tune_trial(ray_start_4_cpus):
+    """Test that setup_mlflow works in a Tune trial."""
+
+    def train_func(config):
+        setup_mlflow(
+            experiment_name="test_exp",
+            create_experiment_if_not_exists=True,
+        )
+
+    tuner = Tuner(train_func)
+    result_grid = tuner.fit()
+
+    assert all(res.error is None for res in result_grid)
+
+
+class MLflowTest(unittest.TestCase):
+    def setUp(self):
+        self.tracking_uri = "sqlite:///" + tempfile.mkdtemp() + "/mlflow.sqlite"
+        self.registry_uri = "sqlite:///" + tempfile.mkdtemp() + "/mlflow.sqlite"
+
+        client = MlflowClient(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        client.create_experiment(name="existing_experiment")
+        # Mlflow > 2 creates a "Default" experiment which has ID 0, so we start our
+        # test with ID 1.
+        assert client.get_experiment_by_name("existing_experiment").experiment_id == "1"
+
+    def tearDown(self) -> None:
+        pass
+
+    def testMlFlowLoggerCallbackConfig(self):
+        # Explicitly pass in all args.
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+            experiment_name="test_exp",
+        )
+        logger.setup()
+        self.assertEqual(
+            logger.mlflow_util._mlflow.get_tracking_uri(), self.tracking_uri
+        )
+        self.assertEqual(
+            logger.mlflow_util._mlflow.get_registry_uri(), self.registry_uri
+        )
+        self.assertListEqual(
+            [e.name for e in logger.mlflow_util._mlflow.search_experiments()],
+            ["test_exp", "existing_experiment", "Default"],
+        )
+        self.assertEqual(logger.mlflow_util.experiment_id, "2")
+
+        # Check if client recognizes already existing experiment.
+        logger = MLflowLoggerCallback(
+            experiment_name="existing_experiment",
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "1")
+
+        # Pass in experiment name as env var.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "test_exp"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "2")
+
+        # Pass in existing experiment name as env var.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "existing_experiment"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "1")
+
+        # Pass in existing experiment id as env var.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "1"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "1")
+
+        # Pass in non existing experiment id as env var.
+        # This should create a new experiment.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "500"
+        with self.assertRaises(ValueError):
+            logger = MLflowLoggerCallback(
+                tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+            )
+            logger.setup()
+
+        # Experiment id env var should take precedence over name env var.
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "test_exp"
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "1"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
+        )
+        logger.setup()
+        self.assertEqual(logger.mlflow_util.experiment_id, "1")
+
+        # Using tags
+        tags = {"user_name": "John", "git_commit_hash": "abc123"}
+        clear_env_vars()
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "test_tags"
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "1"
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri, registry_uri=self.registry_uri, tags=tags
+        )
+        logger.setup()
+        self.assertEqual(logger.tags, tags)
+
+    @patch("ray.air.integrations.mlflow._MLflowLoggerUtil", Mock_MLflowLoggerUtil)
+    def testMlFlowLoggerLogging(self):
+        clear_env_vars()
+        trial_config = {"par1": "a", "par2": "b"}
+        trial = MockTrial(trial_config, "trial1", 0, "artifact")
+
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+            experiment_name="test1",
+            save_artifact=True,
+            tags={"hello": "world"},
+        )
+        logger.setup()
+
+        # Check if run is created with proper tags.
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+        all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=["2"])
+        self.assertEqual(len(all_runs), 1)
+        # all_runs is a pandas dataframe.
+        all_runs = all_runs.to_dict(orient="records")
+        run = logger.mlflow_util._mlflow.get_run(all_runs[0]["run_id"])
+        self.assertDictEqual(
+            run.data.tags,
+            {"hello": "world", "trial_name": "trial1", "mlflow.runName": "trial1"},
+        )
+        self.assertEqual(logger._trial_runs[trial], run.info.run_id)
+        # Params should be logged.
+        self.assertDictEqual(run.data.params, trial_config)
+
+        # When same trial is started again, new run should not be created.
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+        all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=["2"])
+        self.assertEqual(len(all_runs), 1)
+
+        # Check metrics are logged properly.
+        result = {
+            "metric1": 0.8,
+            "metric2": 1,
+            "metric3": None,
+            "training_iteration": 0,
+        }
+        logger.on_trial_result(0, [], trial, result)
+        run = logger.mlflow_util._mlflow.get_run(run_id=run.info.run_id)
+        # metric3 is not logged since it cannot be converted to float.
+        self.assertDictEqual(
+            run.data.metrics, {"metric1": 0.8, "metric2": 1.0, "training_iteration": 0}
+        )
+
+        # Check that artifact is logged on termination.
+        logger.on_trial_complete(0, [], trial)
+        self.assertTrue(logger.mlflow_util.artifact_saved)
+        self.assertDictEqual(
+            logger.mlflow_util.artifact_info,
+            {"dir": "artifact", "run_id": run.info.run_id},
+        )
+
+        # Check if params are logged at the end.
+        run = logger.mlflow_util._mlflow.get_run(run_id=run.info.run_id)
+        self.assertDictEqual(run.data.params, trial_config)
+
+    @patch("ray.air.integrations.mlflow._MLflowLoggerUtil", Mock_MLflowLoggerUtil)
+    def testMlFlowLoggerLogging_logAtEnd(self):
+        clear_env_vars()
+        trial_config = {"par1": "a", "par2": "b"}
+        trial = MockTrial(trial_config, "trial1", 0, "artifact")
+
+        logger = MLflowLoggerCallback(
+            tracking_uri=self.tracking_uri,
+            registry_uri=self.registry_uri,
+            experiment_name="test_log_at_end",
+            tags={"hello": "world"},
+            log_params_on_trial_end=True,
+        )
+        logger.setup()
+        exp_id = logger.mlflow_util.experiment_id
+
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+        all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=[exp_id])
+        self.assertEqual(len(all_runs), 1)
+        # all_runs is a pandas dataframe.
+        all_runs = all_runs.to_dict(orient="records")
+        run = logger.mlflow_util._mlflow.get_run(all_runs[0]["run_id"])
+
+        # Params should NOT be logged at start.
+        self.assertDictEqual(run.data.params, {})
+
+        # Check that params are logged at the end.
+        logger.on_trial_complete(0, [], trial)
+        run = logger.mlflow_util._mlflow.get_run(run_id=run.info.run_id)
+        self.assertDictEqual(run.data.params, trial_config)
+
+    def testMlFlowSetupExplicit(self):
+        clear_env_vars()
+        trial_config = {"par1": 4, "par2": 9.0}
+
+        # No MLflow config passed in.
+        with self.assertRaises(ValueError):
+            setup_mlflow(trial_config)
+
+        # Invalid experiment-id
+        with self.assertRaises(ValueError):
+            setup_mlflow(trial_config, experiment_id="500")
+
+        # Set to experiment that does not already exist.
+        with self.assertRaises(ValueError):
+            setup_mlflow(
+                trial_config,
+                experiment_id="500",
+                experiment_name="new_experiment",
+                tracking_uri=self.tracking_uri,
+            )
+
+        mlflow = setup_mlflow(
+            trial_config,
+            experiment_id="500",
+            experiment_name="existing_experiment",
+            tracking_uri=self.tracking_uri,
+        )
+        mlflow.end_run()
+
+    @patch("ray.train.get_context")
+    def testMlFlowSetupRankNonRankZero(self, mock_get_context):
+        """Assert that non-rank-0 workers get a noop module"""
+        mock_context = MagicMock()
+        mock_context.get_world_rank.return_value = 1
+
+        mock_get_context.return_value = mock_context
+
+        mlflow = setup_mlflow({})
+        assert isinstance(mlflow, _NoopModule)
+
+        mlflow.log_metrics()
+        mlflow.sklearn.save_model(None, "model_directory")
+
+
+class MLflowUtilTest(unittest.TestCase):
+    def setUp(self):
+        self.dirpath = tempfile.mkdtemp()
+        import mlflow
+
+        mlflow.set_tracking_uri("sqlite:///" + self.dirpath + "/mlflow.sqlite")
+        mlflow.create_experiment(name="existing_experiment")
+
+        self.mlflow_util = _MLflowLoggerUtil()
+        self.tracking_uri = mlflow.get_tracking_uri()
+
+    def tearDown(self):
+        shutil.rmtree(self.dirpath)
+
+    def test_experiment_id(self):
+        self.mlflow_util.setup_mlflow(tracking_uri=self.tracking_uri, experiment_id="0")
+        assert self.mlflow_util.experiment_id == "0"
+
+    def test_experiment_id_env_var(self):
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "0"
+        self.mlflow_util.setup_mlflow(tracking_uri=self.tracking_uri)
+        assert self.mlflow_util.experiment_id == "0"
+        del os.environ["MLFLOW_EXPERIMENT_ID"]
+
+    def test_experiment_name(self):
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="existing_experiment"
+        )
+        assert self.mlflow_util.experiment_id == "1"
+
+    def test_run_started_with_correct_experiment(self):
+        experiment_name = "my_experiment_name"
+        # Make sure run is started under the correct experiment.
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name=experiment_name
+        )
+        run = self.mlflow_util.start_run(set_active=True)
+        assert (
+            run.info.experiment_id
+            == self.mlflow_util._mlflow.get_experiment_by_name(
+                experiment_name
+            ).experiment_id
+        )
+
+        self.mlflow_util.end_run()
+
+    def test_experiment_name_env_var(self):
+        os.environ["MLFLOW_EXPERIMENT_NAME"] = "existing_experiment"
+        self.mlflow_util.setup_mlflow(tracking_uri=self.tracking_uri)
+        assert self.mlflow_util.experiment_id == "1"
+        del os.environ["MLFLOW_EXPERIMENT_NAME"]
+
+    def test_id_precedence(self):
+        os.environ["MLFLOW_EXPERIMENT_ID"] = "0"
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="new_experiment"
+        )
+        assert self.mlflow_util.experiment_id == "0"
+        del os.environ["MLFLOW_EXPERIMENT_ID"]
+
+    def test_new_experiment(self):
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="new_experiment"
+        )
+        assert self.mlflow_util.experiment_id == "2"
+
+    def test_setup_fail(self):
+        with self.assertRaises(ValueError):
+            self.mlflow_util.setup_mlflow(
+                tracking_uri=self.tracking_uri,
+                experiment_name="new_experiment2",
+                create_experiment_if_not_exists=False,
+            )
+
+    def test_log_params(self):
+        params = {"a": "a", "x": {"y": "z"}}
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="new_experiment"
+        )
+        run = self.mlflow_util.start_run()
+        run_id = run.info.run_id
+        self.mlflow_util.log_params(params_to_log=params, run_id=run_id)
+
+        run = self.mlflow_util._mlflow.get_run(run_id=run_id)
+        assert run.data.params == flatten_dict(params)
+
+        params2 = {"b": "b"}
+        self.mlflow_util.start_run(set_active=True)
+        self.mlflow_util.log_params(params_to_log=params2, run_id=run_id)
+        run = self.mlflow_util._mlflow.get_run(run_id=run_id)
+        assert run.data.params == flatten_dict(
+            {
+                **params,
+                **params2,
+            }
+        )
+
+        self.mlflow_util.end_run()
+
+    def test_log_metrics(self):
+        metrics = {"a": 1.0, "x": {"y": 2.0}}
+        self.mlflow_util.setup_mlflow(
+            tracking_uri=self.tracking_uri, experiment_name="new_experiment"
+        )
+        run = self.mlflow_util.start_run()
+        run_id = run.info.run_id
+        self.mlflow_util.log_metrics(metrics_to_log=metrics, run_id=run_id, step=0)
+
+        run = self.mlflow_util._mlflow.get_run(run_id=run_id)
+        assert run.data.metrics == flatten_dict(metrics)
+
+        metrics2 = {"b": 1.0}
+        self.mlflow_util.start_run(set_active=True)
+        self.mlflow_util.log_metrics(metrics_to_log=metrics2, run_id=run_id, step=0)
+        assert self.mlflow_util._mlflow.get_run(
+            run_id=run_id
+        ).data.metrics == flatten_dict(
+            {
+                **metrics,
+                **metrics2,
+            }
+        )
+        self.mlflow_util.end_run()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_integration_wandb.py

@@ -0,0 +1,591 @@
+"""Tests for wandb integration.
+
+Note: These tests use a set of mocked APIs:
+- _MockWandbAPI: Mocks wandb API calls (ex: wandb.init).
+- _MockWandbLoggingActor: The same as the regular _WandbLoggingActor,
+    except using the mocked wandb API
+- WandbTestExperimentLogger: Thin subclass of `WandbLoggerCallback` to use for testing.
+    Provides a helper `trial_logging_actors` property that can be used to
+    access attributes of the remote actors for assertions.
+- Use the `get_mock_wandb_logger` helper method to create a logger with
+    a custom mock wandb API class. (Ex: If you want to override some wandb API methods.)
+
+Template for testing with these mocks:
+
+    wandb_logger_kwargs = {}
+    logger = get_mock_wandb_logger(mock_api_cls=_MockWandbAPI, **wandb_logger_kwargs)
+    logger.setup()
+
+    # From now on, the API key is in the env variable.
+    # Start the remote logging actor
+    logger.on_trial_start(0, [], trial)
+    # Log some results
+    result = {}
+    logger.on_trial_result(0, [], trial, result)
+    # Send a STOP signal to the logging actor
+    logger.on_trial_complete(0, [], trial)
+    # This will wait for the logging actor to finish + cleanup
+    logger.on_experiment_end(trials=[trial])
+
+    # Now, we can access properties of the logging actors
+    # (must happen after `on_trial_end` and `on_experiment_end`)
+    logger_state = logger.trial_logging_actor_states[trial]
+    # logger_state.logs, logger_state.config, logger_state.kwargs, ...
+"""
+
+import gc
+import os
+import tempfile
+import time
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+import numpy as np
+import pytest
+
+import ray
+from ray.air.integrations.wandb import (
+    WANDB_ENV_VAR,
+    WANDB_GROUP_ENV_VAR,
+    WANDB_POPULATE_RUN_LOCATION_HOOK,
+    WANDB_PROJECT_ENV_VAR,
+    WANDB_SETUP_API_KEY_HOOK,
+    RunDisabled,
+    WandbLoggerCallback,
+    _QueueItem,
+    _WandbLoggingActor,
+    setup_wandb,
+)
+from ray.air.tests.mocked_wandb_integration import (
+    Trial,
+    WandbTestExperimentLogger,
+    _MockWandbAPI,
+    _MockWandbLoggingActor,
+    get_mock_wandb_logger,
+)
+from ray.exceptions import RayActorError
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+
+
+@pytest.fixture(autouse=True, scope="module")
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def trial():
+    trial_config = {"par1": 4, "par2": 9.12345678}
+    trial = Trial(
+        trial_config,
+        0,
+        "trial_0",
+        "trainable",
+        PlacementGroupFactory([{"CPU": 1}]),
+        "/tmp",
+    )
+    yield trial
+
+
+@pytest.fixture(autouse=True)
+def wandb_env():
+    """Clean up W&B env var before and after each test.
+
+    Even if we use monkeypatch in the test, this is useful to remove environment
+    variables that are set on the laptop when running tests locally.
+    """
+    if WANDB_ENV_VAR in os.environ:
+        del os.environ[WANDB_ENV_VAR]
+    yield
+    if WANDB_ENV_VAR in os.environ:
+        del os.environ[WANDB_ENV_VAR]
+
+
+def fake_wandb_populate_run_location_hook():
+    """Fake user-provided hook to populate W&B environment variables."""
+    os.environ[WANDB_PROJECT_ENV_VAR] = "test_project"
+    os.environ[WANDB_GROUP_ENV_VAR] = "test_group"
+
+
+FAKE_WANDB_POPULATE_RUN_LOCATION_HOOK_IMPORT_PATH = (
+    "ray.air.tests.test_integration_wandb.fake_wandb_populate_run_location_hook"
+)
+
+
+class TestWandbLogger:
+    def test_wandb_logger_project_group(self, monkeypatch):
+        monkeypatch.setenv(WANDB_PROJECT_ENV_VAR, "test_project_from_env_var")
+        monkeypatch.setenv(WANDB_GROUP_ENV_VAR, "test_group_from_env_var")
+        # Read project and group name from environment variable
+        logger = WandbTestExperimentLogger(api_key="1234")
+        logger.setup()
+        assert logger.project == "test_project_from_env_var"
+        assert logger.group == "test_group_from_env_var"
+
+    def test_wandb_logger_api_key_config(self, monkeypatch):
+        # No API key
+        with pytest.raises(ValueError):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+
+        # Fetch API key from argument even if external hook and WANDB_ENV_VAR set
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+        monkeypatch.setenv(
+            WANDB_ENV_VAR,
+            "abcde",
+        )
+        # API Key in config
+        logger = WandbTestExperimentLogger(project="test_project", api_key="1234")
+        logger.setup()
+        assert os.environ[WANDB_ENV_VAR] == "1234"
+
+    def test_wandb_logger_api_key_file(self, monkeypatch):
+        # Fetch API key from file even if external hook and WANDB_ENV_VAR set
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+        monkeypatch.setenv(
+            WANDB_ENV_VAR,
+            "abcde",
+        )
+        # API Key file
+        with tempfile.NamedTemporaryFile("wt") as fp:
+            fp.write("5678")
+            fp.flush()
+
+            logger = WandbTestExperimentLogger(
+                project="test_project", api_key_file=fp.name
+            )
+            logger.setup()
+            assert os.environ[WANDB_ENV_VAR] == "5678"
+
+    def test_wandb_logger_api_key_env_var(self, monkeypatch):
+        # API Key from env var takes precedence over external hook and
+        # logged in W&B API key
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+        monkeypatch.setenv(
+            WANDB_ENV_VAR,
+            "1234",
+        )
+        mock_wandb = Mock(api=Mock(api_key="efgh"))
+        with patch.multiple("ray.air.integrations.wandb", wandb=mock_wandb):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+        assert os.environ[WANDB_ENV_VAR] == "1234"
+        mock_wandb.ensure_configured.assert_not_called()
+
+    def test_wandb_logger_api_key_external_hook(self, monkeypatch):
+        # API Key from external hook if API key not provided through
+        # argument or WANDB_ENV_VAR and user not already logged in to W&B
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+
+        mock_wandb = Mock(api=Mock(api_key=None))
+        with patch.multiple("ray.air.integrations.wandb", wandb=mock_wandb):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+        assert os.environ[WANDB_ENV_VAR] == "abcd"
+        mock_wandb.ensure_configured.assert_called_once()
+
+        mock_wandb = Mock(ensure_configured=Mock(side_effect=AttributeError()))
+        with patch.multiple("ray.air.integrations.wandb", wandb=mock_wandb):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+        assert os.environ[WANDB_ENV_VAR] == "abcd"
+
+    def test_wandb_logger_api_key_from_wandb_login(self, monkeypatch):
+        # No API key should get set if user is already logged in to W&B
+        # and they didn't pass API key through argument or env var.
+        # External hook should not be called because user already logged
+        # in takes precedence.
+        monkeypatch.setenv(
+            WANDB_SETUP_API_KEY_HOOK, "ray._private.test_utils.wandb_setup_api_key_hook"
+        )
+        mock_wandb = Mock()
+        with patch.multiple("ray.air.integrations.wandb", wandb=mock_wandb):
+            logger = WandbTestExperimentLogger(project="test_project")
+            logger.setup()
+        assert os.environ.get(WANDB_ENV_VAR) is None
+        mock_wandb.ensure_configured.assert_called_once()
+
+    def test_wandb_logger_run_location_external_hook(self, monkeypatch):
+        with patch.dict(os.environ):
+            # No project
+            with pytest.raises(ValueError):
+                logger = WandbTestExperimentLogger(api_key="1234")
+                logger.setup()
+
+            # Project and group env vars from external hook
+            monkeypatch.setenv(
+                WANDB_POPULATE_RUN_LOCATION_HOOK,
+                FAKE_WANDB_POPULATE_RUN_LOCATION_HOOK_IMPORT_PATH,
+            )
+            logger = WandbTestExperimentLogger(api_key="1234")
+            logger.setup()
+            assert os.environ[WANDB_PROJECT_ENV_VAR] == "test_project"
+            assert os.environ[WANDB_GROUP_ENV_VAR] == "test_group"
+
+    def test_wandb_logger_start(self, monkeypatch, trial):
+        monkeypatch.setenv(WANDB_ENV_VAR, "9012")
+        # API Key in env
+        logger = WandbTestExperimentLogger(project="test_project")
+        logger.setup()
+        # From now on, the API key is in the env variable.
+        logger.log_trial_start(trial)
+        logger.log_trial_end(trial)
+        logger.on_experiment_end(trials=[trial])
+
+        logger_state = logger.trial_logging_actor_states[trial]
+        assert logger_state.kwargs["project"] == "test_project"
+        assert logger_state.kwargs["id"] == trial.trial_id
+        assert logger_state.kwargs["name"] == trial.trial_name
+        assert logger_state.kwargs["group"] == trial.experiment_dir_name
+        assert "config" in logger_state.exclude
+
+        del logger
+
+        # log config.
+        logger = WandbTestExperimentLogger(project="test_project", log_config=True)
+        logger.log_trial_start(trial)
+        logger.log_trial_end(trial)
+        logger.on_experiment_end(trials=[trial])
+
+        logger_state = logger.trial_logging_actor_states[trial]
+        assert "config" not in logger_state.exclude
+        assert "metric" not in logger_state.exclude
+
+        del logger
+
+        # Exclude metric.
+        logger = WandbTestExperimentLogger(project="test_project", excludes=["metric"])
+        logger.log_trial_start(trial)
+        logger.log_trial_end(trial)
+        logger.on_experiment_end(trials=[trial])
+
+        logger_state = logger.trial_logging_actor_states[trial]
+        assert "config" in logger_state.exclude
+        assert "metric" in logger_state.exclude
+
+        del logger
+
+    def test_wandb_logger_reporting(self, trial):
+        logger = WandbTestExperimentLogger(
+            project="test_project", api_key="1234", excludes=["metric2"]
+        )
+        logger.on_trial_start(0, [], trial)
+        r1 = {
+            "metric1": 0.8,
+            "metric2": 1.4,
+            "metric3": np.asarray(32.0),
+            "metric4": np.float32(32.0),
+            "const": "text",
+            "config": trial.config,
+        }
+        logger.on_trial_result(0, [], trial, r1)
+        logger.on_trial_complete(0, [], trial)
+        logger.on_experiment_end(trials=[trial])
+        logged = logger.trial_logging_actor_states[trial].logs[0]
+        assert "metric1" in logged
+        assert "metric2" not in logged
+        assert "metric3" in logged
+        assert "metric4" in logged
+        assert "const" not in logged
+        assert "config" not in logged
+
+    def test_wandb_logger_auto_config_keys(self, trial):
+        logger = WandbTestExperimentLogger(project="test_project", api_key="1234")
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+        result = {key: 0 for key in WandbLoggerCallback.AUTO_CONFIG_KEYS}
+        logger.on_trial_result(0, [], trial, result)
+        logger.on_trial_complete(0, [], trial)
+        logger.on_experiment_end(trials=[trial])
+        config = logger.trial_logging_actor_states[trial].config
+        # The results in `AUTO_CONFIG_KEYS` should be saved as training configuration
+        # instead of output metrics.
+        assert set(WandbLoggerCallback.AUTO_CONFIG_KEYS) < set(config)
+
+    def test_wandb_logger_exclude_config(self):
+        trial = Trial(
+            config={"param1": 0, "param2": 0},
+            trial_id=0,
+            trial_name="trial_0",
+            experiment_dir_name="trainable",
+            placement_group_factory=PlacementGroupFactory([{"CPU": 1}]),
+            local_path=tempfile.gettempdir(),
+        )
+        logger = WandbTestExperimentLogger(
+            project="test_project",
+            api_key="1234",
+            excludes=(["param2"] + WandbLoggerCallback.AUTO_CONFIG_KEYS),
+        )
+        logger.on_trial_start(iteration=0, trials=[], trial=trial)
+
+        # We need to test that `excludes` also applies to `AUTO_CONFIG_KEYS`.
+        result = {key: 0 for key in WandbLoggerCallback.AUTO_CONFIG_KEYS}
+        logger.on_trial_result(0, [], trial, result)
+        logger.on_trial_complete(0, [], trial)
+        logger.on_experiment_end(trials=[trial])
+
+        config = logger.trial_logging_actor_states[trial].config
+        assert set(config) == {"param1"}
+
+    def test_set_serializability_result(self, trial):
+        """Tests that objects that contain sets can be serialized by wandb."""
+        logger = WandbTestExperimentLogger(
+            project="test_project", api_key="1234", excludes=["metric2"]
+        )
+        logger.on_trial_start(0, [], trial)
+
+        # Testing for https://github.com/ray-project/ray/issues/28541
+        rllib_result = {
+            "env": "simple_spread",
+            "framework": "torch",
+            "num_gpus": 1,
+            "num_workers": 20,
+            "num_envs_per_env_runner": 1,
+            "compress_observations": True,
+            "lambda": 0.99,
+            "train_batch_size": 512,
+            "sgd_minibatch_size": 32,
+            "num_sgd_iter": 5,
+            "batch_mode": "truncate_episodes",
+            "entropy_coeff": 0.01,
+            "lr": 2e-05,
+            "multiagent": {
+                "policies": {"shared_policy"},
+                "policy_mapping_fn": lambda x: x,
+            },
+        }
+        logger.on_trial_result(0, [], trial, rllib_result)
+        logger.on_trial_complete(0, [], trial)
+        logger.on_experiment_end(trials=[trial])
+        logged = logger.trial_logging_actor_states[trial].logs[0]
+        assert logged != "serialization error"
+
+    def test_wandb_logging_actor_api_key(self, trial, monkeypatch):
+        """Tests that the wandb API key get propagated as an environment variable to
+        the remote logging actors."""
+
+        def mock_run(actor_cls):
+            return os.environ.get(WANDB_ENV_VAR)
+
+        monkeypatch.setattr(_MockWandbLoggingActor, "run", mock_run)
+
+        logger = WandbLoggerCallback(
+            project="test_project", api_key="1234", excludes=["metric2"]
+        )
+        logger._logger_actor_cls = _MockWandbLoggingActor
+        logger.setup()
+        logger.log_trial_start(trial)
+        actor_env_var = ray.get(logger._trial_logging_futures[trial])
+        assert actor_env_var == "1234"
+
+    def test_wandb_finish(self, trial, tmp_path):
+        """Test that logging actors are cleaned up upon experiment completion."""
+        marker = tmp_path / "hang_marker"
+        marker.write_text("")
+
+        class HangingFinishMockWandbAPI(_MockWandbAPI):
+            def finish(self):
+                while marker.exists():
+                    time.sleep(0.1)
+
+        logger = get_mock_wandb_logger(
+            mock_api_cls=HangingFinishMockWandbAPI,
+            upload_timeout=1.0,
+        )
+        logger.setup()
+        logger.on_trial_start(0, [], trial)
+        logger.on_trial_complete(0, [], trial)
+        # Signalling stop will not cleanup fully due to the hanging finish
+        assert logger._trial_logging_actors
+        marker.unlink()
+        # wandb.finish has ended -> experiment end hook should cleanup actors fully
+        logger.on_experiment_end(trials=[trial])
+        assert not logger._trial_logging_actors
+
+    def test_wandb_kill_hanging_actor(self, trial):
+        """Test that logging actors are killed if exceeding the upload timeout
+        upon experiment completion."""
+
+        class HangingFinishMockWandbAPI(_MockWandbAPI):
+            def finish(self):
+                time.sleep(5)
+
+        logger = get_mock_wandb_logger(
+            mock_api_cls=HangingFinishMockWandbAPI,
+            upload_timeout=0.1,
+        )
+        logger.setup()
+        logger.on_trial_start(0, [], trial)
+        logger.on_trial_complete(0, [], trial)
+        # Signalling stop will not cleanup fully due to the hanging finish
+        assert logger._trial_logging_actors
+        actor = logger._trial_logging_actors[trial]
+        # Experiment end hook should kill actors since upload_timeout < 5
+        logger.on_experiment_end(trials=[trial])
+        assert not logger._trial_logging_actors
+        gc.collect()
+        with pytest.raises(RayActorError):
+            ray.get(actor.get_state.remote())
+
+    def test_wandb_destructor(self, trial):
+        """Test that the WandbLoggerCallback destructor forcefully cleans up
+        logging actors."""
+
+        class SlowFinishMockWandbAPI(_MockWandbAPI):
+            def finish(self):
+                time.sleep(5)
+
+        logger = get_mock_wandb_logger(
+            mock_api_cls=SlowFinishMockWandbAPI,
+            upload_timeout=1.0,
+        )
+
+        logger.setup()
+        # Triggers logging actor run loop
+        logger.on_trial_start(0, [], trial)
+        actor = logger._trial_logging_actors[trial]
+        del logger
+        gc.collect()
+        with pytest.raises(RayActorError):
+            ray.get(actor.get_state.remote())
+
+    def test_wandb_logging_actor_fault_tolerance(self, trial):
+        """Tests that failing wandb logging actors are restarted"""
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            fail_marker = Path(tempdir) / "fail_marker"
+
+            class _FailingWandbLoggingActor(_MockWandbLoggingActor):
+                def _handle_result(self, result):
+                    if (
+                        result.get("training_iteration") == 3
+                        and not fail_marker.exists()
+                    ):
+                        fail_marker.write_text("Ok")
+                        raise SystemExit
+
+                    return super()._handle_result(result)
+
+            logger = WandbLoggerCallback(
+                project="test_project", api_key="1234", excludes=["metric2"]
+            )
+            logger._logger_actor_cls = _FailingWandbLoggingActor
+            logger.setup()
+            logger.log_trial_start(trial)
+
+            actor = logger._trial_logging_actors[trial]
+            queue = logger._trial_queues[trial]
+
+            logger.log_trial_result(1, trial, result={"training_iteration": 1})
+            logger.log_trial_result(2, trial, result={"training_iteration": 2})
+            logger.log_trial_result(3, trial, result={"training_iteration": 3})
+
+            logger.log_trial_result(4, trial, result={"training_iteration": 4})
+            logger.log_trial_result(5, trial, result={"training_iteration": 5})
+
+            queue.put((_QueueItem.END, None))
+
+            # Wait for the actor's run method to complete
+            ray.get(logger._trial_logging_futures[trial])
+
+            state = ray.get(actor.get_state.remote())
+            assert [metrics["training_iteration"] for metrics in state.logs] == [4, 5]
+
+    def test_wandb_restart(self, trial):
+        """Test that the WandbLoggerCallback reuses actors for trial restarts."""
+
+        logger = WandbLoggerCallback(project="test_project", api_key="1234")
+        logger._logger_actor_cls = _MockWandbLoggingActor
+        logger.setup()
+
+        assert len(logger._trial_logging_futures) == 0
+        assert len(logger._logging_future_to_trial) == 0
+
+        logger.log_trial_start(trial)
+
+        assert len(logger._trial_logging_futures) == 1
+        assert len(logger._logging_future_to_trial) == 1
+
+        logger.log_trial_start(trial)
+
+        assert len(logger._trial_logging_futures) == 1
+        assert len(logger._logging_future_to_trial) == 1
+
+
+def test_wandb_logging_process_run_info_hook(monkeypatch):
+    """
+    Test WANDB_PROCESS_RUN_INFO_HOOK in _WandbLoggingActor is
+    correctly called by calling _WandbLoggingActor.run() mocking
+    out calls to wandb.
+    """
+    mock_queue = Mock(get=Mock(return_value=(_QueueItem.END, None)))
+    monkeypatch.setenv(
+        "WANDB_PROCESS_RUN_INFO_HOOK", "mock_wandb_process_run_info_hook"
+    )
+
+    with patch.object(ray.air.integrations.wandb, "load_class") as mock_load_class:
+        logging_process = _WandbLoggingActor(
+            logdir="/tmp", queue=mock_queue, exclude=[], to_config=[]
+        )
+        logging_process._wandb = Mock()
+        logging_process.run()
+
+    logging_process._wandb.init.assert_called_once()
+    run = logging_process._wandb.init.return_value
+    mock_load_class.assert_called_once_with("mock_wandb_process_run_info_hook")
+    external_hook = mock_load_class.return_value
+    external_hook.assert_called_once_with(run)
+    logging_process._wandb.finish.assert_called_once()
+
+
+def test_wandb_logger_rank_zero_only(trial, monkeypatch):
+    """Test that logging is disabled for non-rank-0 workers when rank_zero_only is True."""
+
+    monkeypatch.setenv(
+        WANDB_ENV_VAR,
+        "abcde",
+    )
+
+    mock_session = Mock()
+    mock_session.experiment_name = "test_project"
+    mock_session.trial_name = "trial_0"
+    mock_session.trial_id = "trial_0"
+
+    # Test case 1: rank_zero_only=True, rank 0
+    mock_session.world_rank = 0
+    with patch("ray.air.integrations.wandb.get_session", return_value=mock_session):
+        run = setup_wandb(project="test_project", rank_zero_only=True, _wandb=Mock())
+        assert not isinstance(run, RunDisabled)
+
+    # Test case 2: rank_zero_only=True, non-rank-0
+    mock_session.world_rank = 1
+    with patch("ray.air.integrations.wandb.get_session", return_value=mock_session):
+        run = setup_wandb(project="test_project", rank_zero_only=True, _wandb=Mock())
+        assert isinstance(run, RunDisabled)
+
+    # Test case 3: rank_zero_only=False, any rank
+    mock_session.world_rank = 1
+    with patch("ray.air.integrations.wandb.get_session", return_value=mock_session):
+        run = setup_wandb(project="test_project", rank_zero_only=False, _wandb=Mock())
+        assert not isinstance(run, RunDisabled)
+
+    # Test case 4: rank_zero_only=True, no session
+    with patch("ray.air.integrations.wandb.get_session", return_value=None):
+        run = setup_wandb(project="test_project", rank_zero_only=True, _wandb=Mock())
+        assert not isinstance(run, RunDisabled)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_keras_callback.py

@@ -0,0 +1,223 @@
+import os
+import sys
+from typing import Dict, Tuple
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+import ray
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.constants import TRAIN_DATASET_KEY
+
+if sys.version_info >= (3, 12):
+    # Tensorflow is not installed for Python 3.12 because of keras compatibility.
+    sys.exit(0)
+else:
+    import tensorflow as tf
+
+    from ray.air.integrations.keras import ReportCheckpointCallback
+    from ray.train.tensorflow import TensorflowTrainer
+
+
+class TestReportCheckpointCallback:
+    @pytest.fixture(name="model")
+    def model_fixture(self):
+        model = tf.keras.Sequential(
+            [tf.keras.layers.InputLayer(input_shape=(1,)), tf.keras.layers.Dense(1)]
+        )
+        model.compile(
+            optimizer="sgd",
+            loss="mean_squared_error",
+            metrics=["accuracy"],
+        )
+        return model
+
+    @patch("ray.train.report")
+    @pytest.mark.parametrize(
+        "metrics, expected_metrics_keys",
+        [
+            (None, {"loss", "accuracy", "val_loss", "val_accuracy"}),
+            ("loss", {"loss"}),
+            (["loss", "accuracy"], {"loss", "accuracy"}),
+            ({"spam": "loss"}, {"spam"}),
+        ],
+    )
+    def test_reported_metrics_contain_expected_keys(
+        self, mock_report, metrics, expected_metrics_keys, model
+    ):
+        # Reported metrics contain different keys depending on the value passed to the
+        # `metrics` parameter. This test varies the value of `metrics` and asserts that
+        # the reported keys are correct.
+        model.fit(
+            x=np.zeros((1, 1)),
+            y=np.zeros((1, 1)),
+            validation_data=(np.zeros((1, 1)), np.zeros((1, 1))),
+            callbacks=[ReportCheckpointCallback(metrics=metrics)],
+        )
+
+        for (metrics,), _ in ray.train.report.call_args_list:
+            assert metrics.keys() == expected_metrics_keys
+
+    @patch("ray.train.report")
+    def test_report_with_default_arguments(self, mock_report, model):
+        # This tests `ReportCheckpointCallback` with default arguments. The test
+        # simulates the end of an epoch, and asserts that a metric and checkpoint are
+        # reported.
+        callback = ReportCheckpointCallback()
+        callback.set_model(model)
+
+        callback.on_epoch_end(0, {"loss": 0})
+
+        assert len(ray.train.report.call_args_list) == 1
+        metrics, checkpoint = self.parse_call(ray.train.report.call_args_list[0])
+        assert metrics == {"loss": 0}
+        assert checkpoint is not None
+
+    @patch("ray.train.report")
+    def test_checkpoint_on_list(self, mock_report, model):
+        # This tests `ReportCheckpointCallback` when `checkpoint_on` is a `list`. The
+        # test simulates each event in `checkpoint_on`, and asserts that a checkpoint
+        # is reported for each event.
+        callback = ReportCheckpointCallback(
+            checkpoint_on=["epoch_end", "train_batch_end"]
+        )
+        callback.model = model
+
+        callback.on_train_batch_end(0, {"loss": 0})
+        callback.on_epoch_end(0, {"loss": 0})
+
+        assert len(ray.train.report.call_args_list) == 2
+        _, first_checkpoint = self.parse_call(ray.train.report.call_args_list[0])
+        assert first_checkpoint is not None
+        _, second_checkpoint = self.parse_call(ray.train.report.call_args_list[0])
+        assert second_checkpoint is not None
+
+    @patch("ray.train.report")
+    def test_report_metrics_on_list(self, mock_report, model):
+        # This tests `ReportCheckpointCallback` when `report_metrics_on` is a `list`.
+        # The test simulates each event in `report_metrics_on`, and asserts that metrics
+        # are reported for each event.
+        callback = ReportCheckpointCallback(
+            report_metrics_on=["epoch_end", "train_batch_end"]
+        )
+        callback.model = model
+
+        callback.on_train_batch_end(0, {"loss": 0})
+        callback.on_epoch_end(0, {"loss": 1})
+
+        assert len(ray.train.report.call_args_list) == 2
+        first_metric, _ = self.parse_call(ray.train.report.call_args_list[0])
+        assert first_metric == {"loss": 0}
+        second_metric, _ = self.parse_call(ray.train.report.call_args_list[1])
+        assert second_metric == {"loss": 1}
+
+    @patch("ray.train.report")
+    def test_report_and_checkpoint_on_different_events(self, mock_report, model):
+        # This tests `ReportCheckpointCallback` when `report_metrics_on` and
+        # `checkpoint_on` are different. The test asserts that:
+        # 1. Checkpoints are reported on `checkpoint_on`
+        # 2. Metrics are reported on `report_metrics_on`
+        # 3. Metrics are reported with checkpoints
+        callback = ReportCheckpointCallback(
+            report_metrics_on="train_batch_end", checkpoint_on="epoch_end"
+        )
+        callback.model = model
+
+        callback.on_train_batch_end(0, {"loss": 0})
+        callback.on_epoch_end(0, {"loss": 1})
+
+        assert len(ray.train.report.call_args_list) == 2
+        first_metric, first_checkpoint = self.parse_call(
+            ray.train.report.call_args_list[0]
+        )
+        assert first_metric == {"loss": 0}
+        assert first_checkpoint is None
+        second_metric, second_checkpoint = self.parse_call(
+            ray.train.report.call_args_list[1]
+        )
+        # We should always include metrics, even if it isn't during one of the events
+        # specified in `report_metrics_on`.
+        assert second_metric == {"loss": 1}
+        assert second_checkpoint is not None
+
+    @patch("ray.train.report")
+    def test_report_delete_tempdir(self, mock_report, model):
+        # This tests `ReportCheckpointCallback`. The test simulates the end of an epoch,
+        # and asserts that the temporary checkpoint directory is deleted afterwards.
+        callback = ReportCheckpointCallback()
+        callback.model = model
+
+        callback.on_epoch_end(0, {"loss": 0})
+
+        assert len(ray.train.report.call_args_list) == 1
+        metrics, checkpoint = self.parse_call(ray.train.report.call_args_list[0])
+        assert metrics == {"loss": 0}
+        assert checkpoint is not None
+        assert checkpoint.path is not None
+        assert not os.path.exists(checkpoint.path)
+
+    def parse_call(self, call) -> Tuple[Dict, train.Checkpoint]:
+        (metrics,), kwargs = call
+        checkpoint = kwargs["checkpoint"]
+        return metrics, checkpoint
+
+
+def get_dataset(a=5, b=10, size=1000):
+    items = [i / size for i in range(size)]
+    dataset = ray.data.from_items([{"x": x, "y": a * x + b} for x in items])
+    return dataset
+
+
+def build_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.InputLayer(input_shape=()),
+            # Add feature dimension, expanding (batch_size,) to (batch_size, 1).
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(10),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_model()
+        multi_worker_model.compile(
+            optimizer=tf.keras.optimizers.SGD(learning_rate=config.get("lr", 1e-3)),
+            loss=tf.keras.losses.mean_squared_error,
+            metrics=[tf.keras.metrics.mean_squared_error],
+        )
+
+    dataset = train.get_dataset_shard("train")
+
+    for _ in range(config.get("epoch", 3)):
+        tf_dataset = dataset.to_tf("x", "y", batch_size=32)
+        multi_worker_model.fit(tf_dataset, callbacks=[ReportCheckpointCallback()])
+
+
+def test_keras_callback_e2e():
+    epochs = 3
+    config = {
+        "epochs": epochs,
+    }
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=2),
+        datasets={TRAIN_DATASET_KEY: get_dataset()},
+    )
+    trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/air/tests/test_new_dataset_config.py

@@ -0,0 +1,394 @@
+import random
+from typing import Optional
+from unittest.mock import MagicMock
+
+import pytest
+
+import ray
+from ray import train
+from ray.data import DataIterator
+from ray.data._internal.execution.interfaces.execution_options import (
+    ExecutionOptions,
+    ExecutionResources,
+)
+from ray.tests.conftest import *  # noqa
+from ray.train import DataConfig, ScalingConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+class TestBasic(DataParallelTrainer):
+    def __init__(
+        self, num_workers: int, expect_ds: bool, expect_sizes: Optional[dict], **kwargs
+    ):
+        def train_loop_per_worker():
+            data_shard = train.get_dataset_shard("train")
+            assert isinstance(data_shard, DataIterator), data_shard
+            for k, v in expect_sizes.items():
+                shard = train.get_dataset_shard(k)
+                if v == -1:
+                    assert shard is None, shard
+                else:
+                    count = 0
+                    for batch in shard.iter_batches():
+                        for arr in batch.values():
+                            count += arr.size
+                    assert count == v, shard
+
+        kwargs.pop("scaling_config", None)
+        super().__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            scaling_config=ScalingConfig(num_workers=num_workers),
+            **kwargs,
+        )
+
+
+def test_basic(ray_start_4_cpus):
+    ds = ray.data.range(10)
+
+    # Single worker basic case.
+    test = TestBasic(
+        1,
+        True,
+        {"train": 10, "test": 10},
+        datasets={"train": ds, "test": ds},
+    )
+    test.fit()
+
+    # Single worker, no test ds.
+    test = TestBasic(1, True, {"train": 10, "test": -1}, datasets={"train": ds})
+    test.fit()
+
+    # Two workers, train and test split.
+    test = TestBasic(
+        2, True, {"train": 5, "test": 5}, datasets={"train": ds, "test": ds}
+    )
+    test.fit()
+
+    # Two workers, both split.
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 5},
+        dataset_config=DataConfig(datasets_to_split=["train", "test"]),
+        datasets={"train": ds, "test": ds},
+    )
+    # Test get config.
+    assert isinstance(test.get_dataset_config(), DataConfig)
+    test.fit()
+
+
+def test_split(ray_start_4_cpus):
+    ds = ray.data.range(10)
+
+    # Split all by default
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 5, "val": 5},
+        datasets={"train": ds, "test": ds, "val": ds},
+    )
+    test.fit()
+
+    # Test flag "all"
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 5},
+        datasets={"train": ds, "test": ds},
+        dataset_config=DataConfig(datasets_to_split="all"),
+    )
+
+    # Test split train only.
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 10},
+        datasets={"train": ds, "test": ds},
+        dataset_config=DataConfig(datasets_to_split=["train"]),
+    )
+    test.fit()
+
+    # Test invalid arguments
+    for datasets_to_split in ["train", ("train"), {}]:
+        with pytest.raises(TypeError, match="`datasets_to_split` should be.*"):
+            test = TestBasic(
+                2,
+                True,
+                {"train": 5, "test": 10},
+                datasets={"train": ds, "test": ds},
+                dataset_config=DataConfig(datasets_to_split=datasets_to_split),
+            )
+
+    # Test empty `datasets_to_split` list
+    test = TestBasic(
+        2,
+        True,
+        {"train": 10, "test": 10},
+        datasets={"train": ds, "test": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+
+def test_configure_execution_options_carryover_context(ray_start_4_cpus):
+    """Tests that execution options in DataContext are carried over to DatConfig
+    automatically."""
+
+    ctx = ray.data.DataContext.get_current()
+    ctx.execution_options.preserve_order = True
+    ctx.execution_options.verbose_progress = True
+
+    data_config = DataConfig()
+
+    ingest_options = data_config.default_ingest_options()
+    assert ingest_options.preserve_order is True
+    assert ingest_options.verbose_progress is True
+
+
+@pytest.mark.parametrize("enable_locality", [True, False])
+def test_configure_locality(enable_locality):
+    data_config = DataConfig(enable_shard_locality=enable_locality)
+
+    mock_ds = MagicMock()
+    mock_ds.streaming_split = MagicMock()
+    mock_ds.copy = MagicMock(return_value=mock_ds)
+    world_size = 2
+    worker_handles = [MagicMock() for _ in range(world_size)]
+    worker_node_ids = ["node" + str(i) for i in range(world_size)]
+    data_config.configure(
+        datasets={"train": mock_ds},
+        world_size=world_size,
+        worker_handles=worker_handles,
+        worker_node_ids=worker_node_ids,
+    )
+    mock_ds.streaming_split.assert_called_once()
+    mock_ds.streaming_split.assert_called_with(
+        world_size,
+        equal=True,
+        locality_hints=worker_node_ids if enable_locality else None,
+    )
+
+
+class CustomConfig(DataConfig):
+    def __init__(self):
+        pass
+
+    def configure(self, *args, **kwargs):
+        ds = ray.data.range(10)
+        return [
+            {"train": ds.iterator()},
+            {"train": ds.iterator()},
+        ]
+
+
+def test_custom_config_subclass(ray_start_4_cpus):
+    test = TestBasic(
+        1,
+        True,
+        {"train": 10},
+        dataset_config=CustomConfig(),
+    )
+    test.fit()
+
+
+class TestRandom(DataParallelTrainer):
+    def __init__(self, num_workers: int, expect_random: bool, **kwargs):
+        def train_loop_per_worker():
+            data_shard = train.get_dataset_shard("train")
+            assert isinstance(data_shard, DataIterator), data_shard
+            epoch1 = list(data_shard.iter_rows())
+            epoch2 = list(data_shard.iter_rows())
+            print("Epochs", epoch1, "\n", epoch2)
+            if expect_random:
+                assert epoch1 != epoch2
+            else:
+                assert epoch1 == epoch2
+
+        kwargs.pop("scaling_config", None)
+        super().__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            scaling_config=ScalingConfig(num_workers=num_workers),
+            **kwargs,
+        )
+
+
+def test_per_epoch_preprocessing(ray_start_4_cpus):
+    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
+    test = TestRandom(2, True, datasets={"train": ds})
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
+    test = TestRandom(2, True, datasets={"train": ds})
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).map(
+        lambda x: {"id": x["id"] * random.random()}
+    )
+    test = TestRandom(2, True, datasets={"train": ds})
+    test.fit()
+
+
+def test_materialized_preprocessing(ray_start_4_cpus):
+    # TODO(ekl) we should test all these configs with splitting enabled, but this
+    # requires implementing deterministic streaming split.
+    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
+    ds = ds.materialize()
+    test = TestRandom(
+        2,
+        False,
+        datasets={"train": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
+    ds = ds.materialize()
+    test = TestRandom(
+        2,
+        False,
+        datasets={"train": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).map(
+        lambda x: {"id": x["id"] * random.random()}
+    )
+    ds = ds.materialize()
+    test = TestRandom(
+        2,
+        False,
+        datasets={"train": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+
+def _run_data_config_resource_test(data_config):
+    cluster_cpus, cluster_gpus = 20, 10
+    num_workers = 2
+    # Resources used by training workers.
+    cpus_per_worker, gpus_per_worker = 2, 1
+
+    original_execution_options = data_config._get_execution_options("train")
+
+    ray.init(num_cpus=cluster_cpus, num_gpus=cluster_gpus)
+
+    class MyTrainer(DataParallelTrainer):
+        def __init__(self, **kwargs):
+            def train_loop_fn():
+                train_ds = train.get_dataset_shard("train")
+                new_execution_options = train_ds.get_context().execution_options
+                if original_execution_options.is_resource_limits_default():
+                    # If the original resource limits are default, the new resource
+                    # limits should be the default as well.
+                    assert new_execution_options.is_resource_limits_default()
+                    exclude_resources = new_execution_options.exclude_resources
+                    assert (
+                        exclude_resources.cpu
+                        == original_execution_options.exclude_resources.cpu
+                        + cpus_per_worker * num_workers
+                        + 1  # trainer coordinator
+                    )
+                    assert (
+                        exclude_resources.gpu
+                        == original_execution_options.exclude_resources.gpu
+                        + gpus_per_worker * num_workers
+                    )
+                else:
+                    # If the original resource limits are not default, the new resource
+                    # limits should be the same as the original ones.
+                    # And the new exclude_resources should be zero.
+                    assert (
+                        new_execution_options.resource_limits
+                        == original_execution_options.resource_limits
+                    )
+                    assert (
+                        new_execution_options.exclude_resources
+                        == ExecutionResources.zero()
+                    )
+
+            kwargs.pop("scaling_config", None)
+
+            super().__init__(
+                train_loop_per_worker=train_loop_fn,
+                scaling_config=ScalingConfig(
+                    num_workers=num_workers,
+                    use_gpu=True,
+                    resources_per_worker={
+                        "CPU": cpus_per_worker,
+                        "GPU": gpus_per_worker,
+                    },
+                ),
+                datasets={"train": ray.data.range(10)},
+                dataset_config=data_config,
+                **kwargs,
+            )
+
+    trainer = MyTrainer()
+    trainer.fit()
+
+
+def test_data_config_default_resource_limits(shutdown_only):
+    """Test that DataConfig preserves user-configured exclude_resources."""
+    execution_options = ExecutionOptions()
+    execution_options.exclude_resources = execution_options.exclude_resources.copy(
+        cpu=2, gpu=1
+    )
+    data_config = DataConfig(execution_options=execution_options)
+
+    _run_data_config_resource_test(data_config)
+
+
+def test_data_config_manual_resource_limits(shutdown_only):
+    """Test manually setting resource limits in DataConfig."""
+    execution_options = ExecutionOptions()
+    execution_options.resource_limits = execution_options.resource_limits.copy(
+        cpu=10, gpu=5
+    )
+    data_config = DataConfig(execution_options=execution_options)
+
+    _run_data_config_resource_test(data_config)
+
+
+def test_v1_train_with_v2_data_autoscaler_sets_exclude_resources(
+    shutdown_only, monkeypatch
+):
+    """Regression test for the Train V1 + V2 cluster autoscaler combination."""
+    monkeypatch.setenv("RAY_DATA_CLUSTER_AUTOSCALER", "V2")
+
+    ray.init(num_cpus=10, num_gpus=2)
+
+    num_train_cpus, num_train_gpus = 4.0, 2.0
+    data_config = DataConfig()
+    data_config.set_train_total_resources(
+        num_train_cpus=num_train_cpus, num_train_gpus=num_train_gpus
+    )
+
+    iterators = data_config.configure(
+        datasets={"train": ray.data.range(10)},
+        world_size=2,
+        worker_handles=None,
+        worker_node_ids=None,
+    )
+
+    exclude_resources = (
+        iterators[0]["train"].get_context().execution_options.exclude_resources
+    )
+    assert exclude_resources.cpu == num_train_cpus
+    assert exclude_resources.gpu == num_train_gpus
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/air/tests/test_remote_storage_hdfs.py

@@ -0,0 +1,50 @@
+"""Test remote_storage in a ci environment with real hdfs setup."""
+
+import os
+
+import pytest
+
+from ray.train.v2._internal.execution.storage import (
+    _list_at_fs_path,
+    _upload_to_fs_path,
+    get_fs_and_path,
+)
+
+
+@pytest.fixture
+def setup_hdfs():
+    """Set env vars required by pyarrow to talk to hdfs correctly.
+
+    Returns hostname and port needed for the hdfs uri."""
+
+    # the following file is written in `install-hdfs.sh`.
+    with open("/tmp/hdfs_env", "r") as f:
+        for line in f.readlines():
+            line = line.rstrip("\n")
+            tokens = line.split("=", maxsplit=1)
+            os.environ[tokens[0]] = tokens[1]
+    import sys
+
+    sys.path.insert(0, os.path.join(os.environ["HADOOP_HOME"], "bin"))
+    hostname = os.getenv("CONTAINER_ID")
+    port = os.getenv("HDFS_PORT")
+    yield hostname, port
+
+
+def test_hdfs(tmp_path, setup_hdfs):
+    pytest.skip("TODO: Fix this test")
+
+    hostname, port = setup_hdfs
+    hdfs_uri = f"hdfs://{hostname}:{port}/test/"
+    fs, path = get_fs_and_path(hdfs_uri)
+
+    dummy_file = tmp_path.joinpath("dummy.txt")
+    dummy_file.write_text("dummy")
+    _upload_to_fs_path(dummy_file, fs, path)
+    assert _list_at_fs_path(fs, path) == ["dummy.txt"]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/tests/test_tracebacks.py

@@ -0,0 +1,102 @@
+import pytest
+from tblib import pickling_support
+
+import ray
+from ray import cloudpickle
+from ray.air._internal.util import StartTraceback, exception_cause, skip_exceptions
+from ray.tune import Tuner
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def _failing_recursive(levels: int = 0, start_traceback: int = -1):
+    if levels > 0:
+        if start_traceback == 0:
+            try:
+                _failing_recursive(
+                    levels=levels - 1, start_traceback=start_traceback - 1
+                )
+            except Exception as e:
+                raise StartTraceback from e
+        else:
+            _failing_recursive(levels=levels - 1, start_traceback=start_traceback - 1)
+    else:
+        raise RuntimeError("Failing")
+
+
+@pytest.mark.parametrize("levels", [4, 5, 6, 7, 8, 9, 10])
+def test_short_traceback(levels):
+    start_traceback = 3
+    with pytest.raises(StartTraceback) as exc_info:
+        _failing_recursive(levels=levels, start_traceback=start_traceback)
+
+    exc = skip_exceptions(exc_info.value)
+    tb = exc.__traceback__
+    i = 0
+    while tb:
+        i += 1
+        tb = tb.tb_next
+
+    assert i == levels - start_traceback + 1
+
+
+def test_recursion():
+    """Test that the skipped exception does not point to the original exception."""
+    root_exception = None
+
+    with pytest.raises(StartTraceback) as exc_info:
+        try:
+            raise Exception("Root Exception")
+        except Exception as e:
+            root_exception = e
+            raise StartTraceback from root_exception
+
+    assert root_exception, "Root exception was not captured."
+
+    start_traceback = exc_info.value
+    skipped_exception = skip_exceptions(start_traceback)
+
+    assert (
+        root_exception != skipped_exception
+    ), "Skipped exception points to the original exception."
+
+
+def test_tblib():
+    """Test that tblib does not cause a maximum recursion error."""
+
+    with pytest.raises(Exception) as exc_info:
+        try:
+            try:
+                raise Exception("Root Exception")
+            except Exception as root_exception:
+                raise StartTraceback from root_exception
+        except Exception as start_traceback:
+            raise skip_exceptions(start_traceback) from exception_cause(start_traceback)
+
+    pickling_support.install()
+    reraised_exception = exc_info.value
+    # This should not raise a RecursionError/PicklingError.
+    cloudpickle.dumps(reraised_exception)
+
+
+def test_traceback_tuner(ray_start_2_cpus):
+    """Ensure that the Tuner's stack trace is not too long."""
+
+    def failing(config):
+        raise RuntimeError("Error")
+
+    tuner = Tuner(failing)
+    results = tuner.fit()
+    assert len(str(results[0].error).split("\n")) <= 20
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/air/tests/test_utils.py

@@ -0,0 +1,34 @@
+"""Test AIR internal utilities (under ray.air._internal)."""
+
+import os
+
+import pytest
+
+import ray
+from ray.air._internal.filelock import RAY_LOCKFILE_DIR, TempFileLock
+
+
+def test_temp_file_lock(tmp_path, monkeypatch):
+    """Test that the directory where temp file locks are saved can be configured
+    via the env variable that configures the global Ray temp dir."""
+    monkeypatch.setenv("RAY_TMPDIR", str(tmp_path))
+    assert str(tmp_path) in ray._common.utils.get_default_system_temp_dir()
+    with TempFileLock(path="abc.txt"):
+        assert RAY_LOCKFILE_DIR in os.listdir(tmp_path)
+        assert os.listdir(tmp_path / RAY_LOCKFILE_DIR)
+
+
+def test_multiple_file_locks(tmp_path, monkeypatch):
+    """Test that a new file lock is created for unique paths."""
+    monkeypatch.setenv("RAY_TMPDIR", str(tmp_path))
+    with TempFileLock(path="abc.txt"):
+        with TempFileLock(path="subdir/abc.txt"):
+            assert RAY_LOCKFILE_DIR in os.listdir(tmp_path)
+            # We should have 2 locks, one for abc.txt and one for subdir/abc.txt
+            assert len(os.listdir(tmp_path / RAY_LOCKFILE_DIR)) == 2
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/air/util/data_batch_conversion.py

@@ -0,0 +1,348 @@
+import warnings
+from enum import Enum
+from typing import TYPE_CHECKING, Dict, List, Union
+
+import numpy as np
+
+from ray.air.constants import TENSOR_COLUMN_NAME
+from ray.air.data_batch_type import DataBatchType
+from ray.data.util.expression_utils import _get_setting_with_copy_warning
+from ray.util.annotations import Deprecated, DeveloperAPI
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+# TODO: Consolidate data conversion edges for arrow bug workaround.
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
+
+# Lazy import to avoid ray init failures without pandas installed and allow
+# dataset to import modules in this file.
+_pandas = None
+
+
+def _lazy_import_pandas():
+    global _pandas
+    if _pandas is None:
+        import pandas
+
+        _pandas = pandas
+    return _pandas
+
+
+@DeveloperAPI
+class BatchFormat(str, Enum):
+    PANDAS = "pandas"
+    # TODO: Remove once Arrow is deprecated as user facing batch format
+    ARROW = "arrow"
+    NUMPY = "numpy"  # Either a single numpy array or a Dict of numpy arrays.
+
+
+@DeveloperAPI
+class BlockFormat(str, Enum):
+    """Internal Dataset block format enum."""
+
+    PANDAS = "pandas"
+    ARROW = "arrow"
+    SIMPLE = "simple"
+
+
+def _convert_batch_type_to_pandas(
+    data: DataBatchType,
+    cast_tensor_columns: bool = False,
+) -> "pd.DataFrame":
+    """Convert the provided data to a Pandas DataFrame.
+
+    Args:
+        data: Data of type DataBatchType
+        cast_tensor_columns: Whether tensor columns should be cast to NumPy ndarrays.
+
+    Returns:
+        A pandas Dataframe representation of the input data.
+
+    """
+    pd = _lazy_import_pandas()
+
+    if isinstance(data, np.ndarray):
+        data = pd.DataFrame({TENSOR_COLUMN_NAME: _ndarray_to_column(data)})
+    elif isinstance(data, dict):
+        tensor_dict = {}
+        for col_name, col in data.items():
+            if not isinstance(col, np.ndarray):
+                raise ValueError(
+                    "All values in the provided dict must be of type "
+                    f"np.ndarray. Found type {type(col)} for key {col_name} "
+                    f"instead."
+                )
+            tensor_dict[col_name] = _ndarray_to_column(col)
+        data = pd.DataFrame(tensor_dict)
+    elif pyarrow is not None and isinstance(data, pyarrow.Table):
+        data = data.to_pandas()
+    elif not isinstance(data, pd.DataFrame):
+        raise ValueError(
+            f"Received data of type: {type(data)}, but expected it to be one "
+            f"of {DataBatchType}"
+        )
+    if cast_tensor_columns:
+        data = _cast_tensor_columns_to_ndarrays(data)
+    return data
+
+
+def _convert_pandas_to_batch_type(
+    data: "pd.DataFrame",
+    type: BatchFormat,
+    cast_tensor_columns: bool = False,
+) -> DataBatchType:
+    """Convert the provided Pandas dataframe to the provided ``type``.
+
+    Args:
+        data: A Pandas DataFrame
+        type: The specific ``BatchFormat`` to convert to.
+        cast_tensor_columns: Whether tensor columns should be cast to our tensor
+            extension type.
+
+    Returns:
+        The input data represented with the provided type.
+    """
+    if cast_tensor_columns:
+        data = _cast_ndarray_columns_to_tensor_extension(data)
+    if type == BatchFormat.PANDAS:
+        return data
+
+    elif type == BatchFormat.NUMPY:
+        if len(data.columns) == 1:
+            # If just a single column, return as a single numpy array.
+            return data.iloc[:, 0].to_numpy()
+        else:
+            # Else return as a dict of numpy arrays.
+            output_dict = {}
+            for column in data:
+                output_dict[column] = data[column].to_numpy()
+            return output_dict
+
+    elif type == BatchFormat.ARROW:
+        if not pyarrow:
+            raise ValueError(
+                "Attempted to convert data to Pyarrow Table but Pyarrow "
+                "is not installed. Please do `pip install pyarrow` to "
+                "install Pyarrow."
+            )
+        return pyarrow.Table.from_pandas(data)
+
+    else:
+        raise ValueError(
+            f"Received type {type}, but expected it to be one of {DataBatchType}"
+        )
+
+
+@Deprecated
+def convert_batch_type_to_pandas(
+    data: DataBatchType,
+    cast_tensor_columns: bool = False,
+):
+    """Convert the provided data to a Pandas DataFrame.
+
+    This API is deprecated from Ray 2.4.
+
+    Args:
+        data: Data of type DataBatchType
+        cast_tensor_columns: Whether tensor columns should be cast to NumPy ndarrays.
+
+    Returns:
+        A pandas Dataframe representation of the input data.
+
+    """
+    warnings.warn(
+        "`convert_batch_type_to_pandas` is deprecated as a developer API "
+        "starting from Ray 2.4. All batch format conversions should be "
+        "done manually instead of relying on this API.",
+        PendingDeprecationWarning,
+    )
+    return _convert_batch_type_to_pandas(
+        data=data, cast_tensor_columns=cast_tensor_columns
+    )
+
+
+@Deprecated
+def convert_pandas_to_batch_type(
+    data: "pd.DataFrame",
+    type: BatchFormat,
+    cast_tensor_columns: bool = False,
+):
+    """Convert the provided Pandas dataframe to the provided ``type``.
+
+    Args:
+        data: A Pandas DataFrame
+        type: The specific ``BatchFormat`` to convert to.
+        cast_tensor_columns: Whether tensor columns should be cast to our tensor
+            extension type.
+
+    Returns:
+        The input data represented with the provided type.
+    """
+    warnings.warn(
+        "`convert_pandas_to_batch_type` is deprecated as a developer API "
+        "starting from Ray 2.4. All batch format conversions should be "
+        "done manually instead of relying on this API.",
+        PendingDeprecationWarning,
+    )
+    return _convert_pandas_to_batch_type(
+        data=data, type=type, cast_tensor_columns=cast_tensor_columns
+    )
+
+
+def _convert_batch_type_to_numpy(
+    data: DataBatchType,
+) -> Union[np.ndarray, Dict[str, np.ndarray]]:
+    """Convert the provided data to a NumPy ndarray or dict of ndarrays.
+
+    Args:
+        data: Data of type DataBatchType
+
+    Returns:
+        A numpy representation of the input data.
+    """
+    pd = _lazy_import_pandas()
+
+    if isinstance(data, np.ndarray):
+        return data
+    elif isinstance(data, dict):
+        for col_name, col in data.items():
+            if not isinstance(col, np.ndarray):
+                raise ValueError(
+                    "All values in the provided dict must be of type "
+                    f"np.ndarray. Found type {type(col)} for key {col_name} "
+                    f"instead."
+                )
+        return data
+    elif pyarrow is not None and isinstance(data, pyarrow.Table):
+        from ray.data._internal.arrow_ops import transform_pyarrow
+        from ray.data._internal.tensor_extensions.arrow import (
+            get_arrow_extension_fixed_shape_tensor_types,
+        )
+
+        column_values_ndarrays = []
+
+        for col in data.columns:
+            # Combine columnar values arrays to make these contiguous
+            # (making them compatible with numpy format)
+            combined_array = transform_pyarrow.combine_chunked_array(col)
+
+            column_values_ndarrays.append(
+                transform_pyarrow.to_numpy(combined_array, zero_copy_only=False)
+            )
+
+        arrow_fixed_shape_tensor_types = get_arrow_extension_fixed_shape_tensor_types()
+
+        # NOTE: This branch is here for backwards-compatibility
+        if data.column_names == [TENSOR_COLUMN_NAME] and (
+            isinstance(data.schema.types[0], arrow_fixed_shape_tensor_types)
+        ):
+            return column_values_ndarrays[0]
+
+        return dict(zip(data.column_names, column_values_ndarrays))
+    elif isinstance(data, pd.DataFrame):
+        return _convert_pandas_to_batch_type(data, BatchFormat.NUMPY)
+    else:
+        raise ValueError(
+            f"Received data of type: {type(data)}, but expected it to be one "
+            f"of {DataBatchType}"
+        )
+
+
+def _ndarray_to_column(arr: np.ndarray) -> Union["pd.Series", List[np.ndarray]]:
+    """Convert a NumPy ndarray into an appropriate column format for insertion into a
+    pandas DataFrame.
+
+    If conversion to a pandas Series fails (e.g. if the ndarray is multi-dimensional),
+    fall back to a list of NumPy ndarrays.
+    """
+    pd = _lazy_import_pandas()
+    try:
+        # Try to convert to Series, falling back to a list conversion if this fails
+        # (e.g. if the ndarray is multi-dimensional).
+        return pd.Series(arr)
+    except ValueError:
+        return list(arr)
+
+
+def _unwrap_ndarray_object_type_if_needed(arr: np.ndarray) -> np.ndarray:
+    """Unwrap an object-dtyped NumPy ndarray containing ndarray pointers into a single
+    contiguous ndarray, if needed/possible.
+    """
+    if arr.dtype.type is np.object_:
+        try:
+            # Try to convert the NumPy ndarray to a non-object dtype.
+            arr = np.array([np.asarray(v) for v in arr])
+        except Exception:
+            # This may fail if the subndarrays are of heterogeneous shape
+            pass
+    return arr
+
+
+def _cast_ndarray_columns_to_tensor_extension(df: "pd.DataFrame") -> "pd.DataFrame":
+    """
+    Cast all NumPy ndarray columns in df to our tensor extension type, TensorArray.
+    """
+    # Get the SettingWithCopyWarning class if available
+    SettingWithCopyWarning = _get_setting_with_copy_warning()
+
+    from ray.data._internal.tensor_extensions.pandas import (
+        TensorArray,
+        column_needs_tensor_extension,
+    )
+
+    # Try to convert any ndarray columns to TensorArray columns.
+    # TODO(Clark): Once Pandas supports registering extension types for type
+    # inference on construction, implement as much for NumPy ndarrays and remove
+    # this. See https://github.com/pandas-dev/pandas/issues/41848
+    # TODO(Clark): Optimize this with propagated DataFrame metadata containing a list of
+    # column names containing tensor columns, to make this an O(# of tensor columns)
+    # check rather than the current O(# of columns) check.
+    for col_name, col in df.items():
+        if column_needs_tensor_extension(col):
+            try:
+                # Suppress Pandas warnings:
+                # https://github.com/ray-project/ray/issues/29270
+                # We actually want in-place operations so we surpress this warning.
+                # https://stackoverflow.com/a/74193599
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", category=FutureWarning)
+                    if SettingWithCopyWarning is not None:
+                        warnings.simplefilter("ignore", category=SettingWithCopyWarning)
+                    df[col_name] = TensorArray(col)
+            except Exception as e:
+                raise ValueError(
+                    f"Tried to cast column {col_name} to the TensorArray tensor "
+                    "extension type but the conversion failed. To disable "
+                    "automatic casting to this tensor extension, set "
+                    "ctx = DataContext.get_current(); "
+                    "ctx.enable_tensor_extension_casting = False."
+                ) from e
+    return df
+
+
+def _cast_tensor_columns_to_ndarrays(df: "pd.DataFrame") -> "pd.DataFrame":
+    """Cast all tensor extension columns in df to NumPy ndarrays."""
+    # Get the SettingWithCopyWarning class if available
+    SettingWithCopyWarning = _get_setting_with_copy_warning()
+    from ray.data._internal.tensor_extensions.pandas import TensorDtype
+
+    # Try to convert any tensor extension columns to ndarray columns.
+    # TODO(Clark): Optimize this with propagated DataFrame metadata containing a list of
+    # column names containing tensor columns, to make this an O(# of tensor columns)
+    # check rather than the current O(# of columns) check.
+    for col_name, col in df.items():
+        if isinstance(col.dtype, TensorDtype):
+            # Suppress Pandas warnings:
+            # https://github.com/ray-project/ray/issues/29270
+            # We actually want in-place operations so we surpress this warning.
+            # https://stackoverflow.com/a/74193599
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", category=FutureWarning)
+                if SettingWithCopyWarning is not None:
+                    warnings.simplefilter("ignore", category=SettingWithCopyWarning)
+                df[col_name] = list(col.to_numpy())
+    return df

python/ray/air/util/node.py

@@ -0,0 +1,65 @@
+from typing import Dict, Optional, Union
+
+import ray
+
+
+def _get_node_id_from_node_ip(node_ip: str) -> Optional[str]:
+    """Returns the node ID for the first alive node with the input IP."""
+    for node in ray.nodes():
+        if node["Alive"] and node["NodeManagerAddress"] == node_ip:
+            return node["NodeID"]
+
+    return None
+
+
+def _force_on_node(
+    node_id: str,
+    remote_func_or_actor_class: Optional[
+        Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass]
+    ] = None,
+) -> Union[Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass], Dict]:
+    """Schedule a remote function or actor class on a given node.
+
+    Args:
+        node_id: The node to schedule on.
+        remote_func_or_actor_class: A Ray remote function or actor class
+            to schedule on the input node. If None, this function will directly
+            return the options dict to pass to another remote function or actor class
+            as remote options.
+    Returns:
+        The provided remote function or actor class, but with options modified to force
+        placement on the input node. If remote_func_or_actor_class is None,
+        the options dict to pass to another remote function or
+        actor class as remote options kwargs.
+    """
+
+    options = {"label_selector": {ray._raylet.RAY_NODE_ID_KEY: node_id}}
+
+    if remote_func_or_actor_class is None:
+        return options
+
+    return remote_func_or_actor_class.options(**options)
+
+
+def _force_on_current_node(
+    remote_func_or_actor_class: Optional[
+        Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass]
+    ] = None
+) -> Union[Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass], Dict]:
+    """Schedule a remote function or actor class on the current node.
+
+    If using Ray Client, the current node is the client server node.
+
+    Args:
+        remote_func_or_actor_class: A Ray remote function or actor class
+            to schedule on the current node. If None, this function will directly
+            return the options dict to pass to another remote function or actor class
+            as remote options.
+    Returns:
+        The provided remote function or actor class, but with options modified to force
+        placement on the current node. If remote_func_or_actor_class is None,
+        the options dict to pass to another remote function or
+        actor class as remote options kwargs.
+    """
+    current_node_id = ray.get_runtime_context().get_node_id()
+    return _force_on_node(current_node_id, remote_func_or_actor_class)

python/ray/air/util/object_extensions/arrow.py

@@ -0,0 +1,7 @@
+# NOTE: We provide these as aliases to maintain compatibility with older version
+#       of Arrow `PyExtensionType` that relies on picked class references that
+#       reference `ray.air.util.{tensor|object}_extensions.arrow.*` classes
+
+from ray.data._internal.object_extensions.arrow import (
+    ArrowPythonObjectType,  # noqa: F401
+)

python/ray/air/util/tensor_extensions/arrow.py

@@ -0,0 +1,9 @@
+# NOTE: We provide these as aliases to maintain compatibility with older version
+#       of Arrow `PyExtensionType` that relies on picked class references that
+#       reference `ray.air.util.tensor_extensions.arrow.*` classes
+
+from ray.data._internal.tensor_extensions.arrow import (
+    ArrowTensorType,  # noqa: F401
+    ArrowTensorTypeV2,  # noqa: F401
+    ArrowVariableShapedTensorType,  # noqa: F401
+)