chore: import upstream snapshot with attribution

python/ray/train/.claude/CLAUDE.md

@@ -0,0 +1,10 @@
+<!-- Loaded on-demand when Claude works on Ray Train files. -->
+<!-- Keep under 50 lines. Multi-step procedures → skills. Code style → rules/. -->
+
+# Ray Train
+
+## Key Modules
+<!-- Entry points, important abstractions, non-obvious dependencies -->
+
+## Gotchas
+<!-- Non-obvious behaviors, common mistakes, things that break silently -->

python/ray/train/.claude/rules/README.md

@@ -0,0 +1,9 @@
+<!-- Add Ray Train team-specific rules here as .md files. -->
+<!-- Rules with paths: frontmatter only load when matching files are edited. -->
+<!-- Example:
+---
+paths:
+  - "python/ray/train/**/*.py"
+---
+- Use the base trainer interface for new training backends
+-->

python/ray/train/__init__.py

@@ -0,0 +1,143 @@
+# Try import ray[train] core requirements (defined in setup.py)
+# isort: off
+try:
+    import fsspec  # noqa: F401
+    import pandas  # noqa: F401
+    import pyarrow  # noqa: F401
+    import requests  # noqa: F401
+except ImportError as exc:
+    raise ImportError(
+        "Can't import ray.train as some dependencies are missing. "
+        'Run `pip install "ray[train]"` to fix.'
+    ) from exc
+# isort: on
+
+
+from ray.air.config import CheckpointConfig, FailureConfig, RunConfig, ScalingConfig
+from ray.air.result import Result
+
+# Import this first so it can be used in other modules
+from ray.train._checkpoint import Checkpoint
+from ray.train._internal.data_config import DataConfig
+from ray.train._internal.session import get_checkpoint, get_dataset_shard, report
+from ray.train._internal.syncer import SyncConfig
+from ray.train.backend import BackendConfig
+from ray.train.base_trainer import TrainingFailedError
+from ray.train.constants import TRAIN_DATASET_KEY
+from ray.train.context import TrainContext, get_context
+from ray.train.v2._internal.constants import is_v2_enabled
+
+if is_v2_enabled():
+    try:
+        import pydantic  # noqa: F401
+    except (ImportError, ModuleNotFoundError) as exc:
+        raise ImportError(
+            "`ray.train.v2` requires the pydantic package, which is missing. "
+            "Run the following command to fix this: `pip install pydantic`"
+        ) from exc
+    from ray.train.v2.api.callback import UserCallback  # noqa: F811
+    from ray.train.v2.api.config import (  # noqa: F811
+        CheckpointConfig,
+        FailureConfig,
+        LoggingConfig,
+        RunConfig,
+        ScalingConfig,
+    )
+    from ray.train.v2.api.context import TrainContext  # noqa: F811
+    from ray.train.v2.api.exceptions import (  # noqa: F811
+        ControllerError,
+        TrainingFailedError,
+        WorkerGroupError,
+    )
+    from ray.train.v2.api.report_config import (  # noqa: F811
+        CheckpointConsistencyMode,
+        CheckpointUploadMode,
+    )
+    from ray.train.v2.api.reported_checkpoint import (  # noqa: F811
+        ReportedCheckpoint,
+        ReportedCheckpointStatus,
+    )
+    from ray.train.v2.api.result import Result  # noqa: F811
+    from ray.train.v2.api.train_fn_utils import (  # noqa: F811
+        get_all_reported_checkpoints,
+        get_checkpoint,
+        get_context,
+        get_dataset_shard,
+        report,
+    )
+    from ray.train.v2.api.validation_config import (  # noqa: F811
+        ValidationConfig,
+        ValidationFn,
+        ValidationTaskConfig,
+    )
+
+
+__all__ = [
+    "get_checkpoint",
+    "get_context",
+    "get_dataset_shard",
+    "report",
+    "BackendConfig",
+    "Checkpoint",
+    "CheckpointConfig",
+    "DataConfig",
+    "FailureConfig",
+    "Result",
+    "RunConfig",
+    "ScalingConfig",
+    "SyncConfig",
+    "TrainContext",
+    "TrainingFailedError",
+    "TRAIN_DATASET_KEY",
+]
+
+get_checkpoint.__module__ = "ray.train"
+get_context.__module__ = "ray.train"
+get_dataset_shard.__module__ = "ray.train"
+report.__module__ = "ray.train"
+BackendConfig.__module__ = "ray.train"
+Checkpoint.__module__ = "ray.train"
+CheckpointConfig.__module__ = "ray.train"
+DataConfig.__module__ = "ray.train"
+FailureConfig.__module__ = "ray.train"
+Result.__module__ = "ray.train"
+RunConfig.__module__ = "ray.train"
+ScalingConfig.__module__ = "ray.train"
+SyncConfig.__module__ = "ray.train"
+TrainContext.__module__ = "ray.train"
+TrainingFailedError.__module__ = "ray.train"
+
+# TODO: consider implementing these in v1 and raising ImportError instead.
+if is_v2_enabled():
+    __all__.extend(
+        [
+            "CheckpointUploadMode",
+            "CheckpointConsistencyMode",
+            "ControllerError",
+            "LoggingConfig",
+            "ReportedCheckpoint",
+            "ReportedCheckpointStatus",
+            "UserCallback",
+            "WorkerGroupError",
+            "ValidationConfig",
+            "ValidationFn",
+            "ValidationTaskConfig",
+            "get_all_reported_checkpoints",
+        ]
+    )
+
+    CheckpointUploadMode.__module__ = "ray.train"
+    CheckpointConsistencyMode.__module__ = "ray.train"
+    ControllerError.__module__ = "ray.train"
+    LoggingConfig.__module__ = "ray.train"
+    ReportedCheckpoint.__module__ = "ray.train"
+    ReportedCheckpointStatus.__module__ = "ray.train"
+    UserCallback.__module__ = "ray.train"
+    WorkerGroupError.__module__ = "ray.train"
+    ValidationConfig.__module__ = "ray.train"
+    ValidationFn.__module__ = "ray.train"
+    ValidationTaskConfig.__module__ = "ray.train"
+    get_all_reported_checkpoints.__module__ = "ray.train"
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/_checkpoint.py

@@ -0,0 +1,434 @@
+import contextlib
+import glob
+import json
+import logging
+import os
+import platform
+import shutil
+import tempfile
+import traceback
+import uuid
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+import pyarrow.fs
+
+from ray.air._internal.filelock import TempFileLock
+from ray.train._internal.storage import _download_from_fs_path, _exists_at_fs_path
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+# The filename of the file that stores user metadata set on the checkpoint.
+_METADATA_FILE_NAME = ".metadata.json"
+
+# The prefix of the temp checkpoint directory that `to_directory` downloads to
+# on the local filesystem.
+_CHECKPOINT_TEMP_DIR_PREFIX = "checkpoint_tmp_"
+
+
+class _CheckpointMetaClass(type):
+    def __getattr__(self, item):
+        try:
+            return super().__getattribute__(item)
+        except AttributeError as exc:
+            if item in {
+                "from_dict",
+                "to_dict",
+                "from_bytes",
+                "to_bytes",
+                "get_internal_representation",
+            }:
+                raise _get_migration_error(item) from exc
+            elif item in {
+                "from_uri",
+                "to_uri",
+                "uri",
+            }:
+                raise _get_uri_error(item) from exc
+            elif item in {"get_preprocessor", "set_preprocessor"}:
+                raise _get_preprocessor_error(item) from exc
+
+            raise exc
+
+
+@PublicAPI(stability="beta")
+class Checkpoint(metaclass=_CheckpointMetaClass):
+    """A reference to data persisted as a directory in local or remote storage.
+
+    Access the checkpoint contents locally using ``checkpoint.to_directory()``
+    or ``checkpoint.as_directory``.
+
+    Attributes
+    ----------
+    path: A path on the filesystem containing the checkpoint contents.
+    filesystem: PyArrow FileSystem that can be used to access data at the `path`.
+
+    See Also
+    --------
+    ray.train.report : Report a checkpoint during training (with Ray Train/Tune).
+    ray.train.get_checkpoint : Get the latest checkpoint during training
+        (for restoration).
+
+    :ref:`train-checkpointing`
+    :ref:`persistent-storage-guide`
+
+    Examples
+    --------
+
+    Creating a checkpoint using ``Checkpoint.from_directory``:
+
+        >>> from ray.train import Checkpoint
+        >>> checkpoint = Checkpoint.from_directory("/tmp/example_checkpoint_dir")
+        >>> checkpoint.filesystem  # doctest: +ELLIPSIS
+        <pyarrow._fs.LocalFileSystem object...
+        >>> checkpoint.path
+        '/tmp/example_checkpoint_dir'
+
+    Creating a checkpoint from a remote URI:
+
+        >>> checkpoint = Checkpoint("s3://bucket/path/to/checkpoint")
+        >>> checkpoint.filesystem  # doctest: +ELLIPSIS
+        <pyarrow._s3fs.S3FileSystem object...
+        >>> checkpoint.path
+        'bucket/path/to/checkpoint'
+
+    Creating a checkpoint with a custom filesystem:
+
+        >>> checkpoint = Checkpoint(
+        ...     path="bucket/path/to/checkpoint",
+        ...     filesystem=pyarrow.fs.S3FileSystem(),
+        ... )
+        >>> checkpoint.filesystem  # doctest: +ELLIPSIS
+        <pyarrow._s3fs.S3FileSystem object...
+        >>> checkpoint.path
+        'bucket/path/to/checkpoint'
+
+    Accessing a checkpoint's contents:
+
+        >>> import os  # doctest: +SKIP
+        >>> with checkpoint.as_directory() as local_checkpoint_dir:  # doctest: +SKIP
+        ...    print(os.listdir(local_checkpoint_dir))  # doctest: +SKIP
+        ['model.pt', 'optimizer.pt', 'misc.pt']
+    """
+
+    def __init__(
+        self,
+        path: Union[str, os.PathLike],
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+    ):
+        """Construct a Checkpoint.
+
+        Args:
+            path: A local path or remote URI containing the checkpoint data.
+                If a filesystem is provided, then this path must NOT be a URI.
+                It should be a path on the filesystem with the prefix already stripped.
+            filesystem: PyArrow FileSystem to use to access data at the path.
+                If not specified, this is inferred from the URI scheme.
+        """
+        self.path = str(path)
+        self.filesystem = filesystem
+
+        if path and not filesystem:
+            self.filesystem, self.path = pyarrow.fs.FileSystem.from_uri(path)
+
+        # This random UUID is used to create a temporary directory name on the
+        # local filesystem, which will be used for downloading checkpoint data.
+        # This ensures that if multiple processes download the same checkpoint object
+        # only one process performs the actual download while the others wait.
+        # This prevents duplicated download efforts and data.
+        # NOTE: Calling `to_directory` from multiple `Checkpoint` objects
+        # that point to the same (fs, path) will still download the data multiple times.
+        # This only ensures a canonical temp directory name for a single `Checkpoint`.
+        self._uuid = uuid.uuid4()
+
+    def __repr__(self):
+        return f"Checkpoint(filesystem={self.filesystem.type_name}, path={self.path})"
+
+    def get_metadata(self) -> Dict[str, Any]:
+        """Return the metadata dict stored with the checkpoint.
+
+        If no metadata is stored, an empty dict is returned.
+        """
+        metadata_path = Path(self.path, _METADATA_FILE_NAME).as_posix()
+        if not _exists_at_fs_path(self.filesystem, metadata_path):
+            return {}
+
+        with self.filesystem.open_input_file(metadata_path) as f:
+            return json.loads(f.readall().decode("utf-8"))
+
+    def set_metadata(self, metadata: Dict[str, Any]) -> None:
+        """Set the metadata stored with this checkpoint.
+
+        This will overwrite any existing metadata stored with this checkpoint.
+        """
+        metadata_path = Path(self.path, _METADATA_FILE_NAME).as_posix()
+        with self.filesystem.open_output_stream(metadata_path) as f:
+            f.write(json.dumps(metadata).encode("utf-8"))
+
+    def update_metadata(self, metadata: Dict[str, Any]) -> None:
+        """Update the metadata stored with this checkpoint.
+
+        This will update any existing metadata stored with this checkpoint.
+        """
+        existing_metadata = self.get_metadata()
+        existing_metadata.update(metadata)
+        self.set_metadata(existing_metadata)
+
+    @classmethod
+    def from_directory(cls, path: Union[str, os.PathLike]) -> "Checkpoint":
+        """Create checkpoint object from a local directory.
+
+        Args:
+            path: Local directory containing checkpoint data.
+
+        Returns:
+            A ray.train.Checkpoint object.
+        """
+        return cls(path, filesystem=pyarrow.fs.LocalFileSystem())
+
+    def to_directory(self, path: Optional[Union[str, os.PathLike]] = None) -> str:
+        """Write checkpoint data to a local directory.
+
+        *If multiple processes on the same node call this method simultaneously,*
+        only a single process will perform the download, while the others
+        wait for the download to finish. Once the download finishes, all processes
+        receive the same local directory to read from.
+
+        Args:
+            path: Target directory to download data to. If not specified,
+                this method will use a temporary directory.
+
+        Returns:
+            str: Directory containing checkpoint data.
+        """
+        user_provided_path = path is not None
+        local_path = (
+            path if user_provided_path else self._get_temporary_checkpoint_dir()
+        )
+        local_path = os.path.normpath(os.path.expanduser(str(local_path)))
+        os.makedirs(local_path, exist_ok=True)
+
+        try:
+            # Timeout 0 means there will be only one attempt to acquire
+            # the file lock. If it cannot be acquired, throw a TimeoutError
+            with TempFileLock(local_path, timeout=0):
+                _download_from_fs_path(
+                    fs=self.filesystem, fs_path=self.path, local_path=local_path
+                )
+        except TimeoutError:
+            # if the directory is already locked, then wait but do not do anything.
+            with TempFileLock(local_path, timeout=-1):
+                pass
+            if not os.path.exists(local_path):
+                raise RuntimeError(
+                    f"Checkpoint directory {local_path} does not exist, "
+                    "even though it should have been created by "
+                    "another process. Please raise an issue on GitHub: "
+                    "https://github.com/ray-project/ray/issues"
+                )
+
+        return local_path
+
+    @contextlib.contextmanager
+    def as_directory(self) -> Iterator[str]:
+        """Returns checkpoint contents in a local directory as a context.
+
+        This function makes checkpoint data available as a directory while avoiding
+        unnecessary copies and left-over temporary data.
+
+        *If the checkpoint points to a local directory*, this method just returns the
+        local directory path without making a copy, and nothing will be cleaned up
+        after exiting the context.
+
+        *If the checkpoint points to a remote directory*, this method will download the
+        checkpoint to a local temporary directory and return the path
+        to the temporary directory.
+
+        *If multiple processes on the same node call this method simultaneously,*
+        only a single process will perform the download, while the others
+        wait for the download to finish. Once the download finishes, all processes
+        receive the same local (temporary) directory to read from.
+
+        Once all processes have finished working with the checkpoint,
+        the temporary directory is cleaned up.
+
+        Users should treat the returned checkpoint directory as read-only and avoid
+        changing any data within it, as it may be deleted when exiting the context.
+
+        Example:
+
+        .. testcode::
+            :hide:
+
+            from pathlib import Path
+            import tempfile
+
+            from ray.train import Checkpoint
+
+            temp_dir = tempfile.mkdtemp()
+            (Path(temp_dir) / "example.txt").write_text("example checkpoint data")
+            checkpoint = Checkpoint.from_directory(temp_dir)
+
+        .. testcode::
+
+            with checkpoint.as_directory() as checkpoint_dir:
+                # Do some read-only processing of files within checkpoint_dir
+                pass
+
+            # At this point, if a temporary directory was created, it will have
+            # been deleted.
+
+        Yields:
+            str: The local path to the checkpoint directory.
+        """
+        if isinstance(self.filesystem, pyarrow.fs.LocalFileSystem):
+            yield self.path
+        else:
+            del_lock_path = _get_del_lock_path(self._get_temporary_checkpoint_dir())
+            open(del_lock_path, "a").close()
+
+            temp_dir = self.to_directory()
+            try:
+                yield temp_dir
+            finally:
+                # Always cleanup the del lock after we're done with the directory.
+                # This avoids leaving a lock file behind in the case of an exception
+                # in the user code.
+                try:
+                    os.remove(del_lock_path)
+                except Exception:
+                    logger.warning(
+                        f"Could not remove {del_lock_path} deletion file lock. "
+                        f"Traceback:\n{traceback.format_exc()}"
+                    )
+
+                # If there are no more lock files, that means there are no more
+                # readers of this directory, and we can safely delete it.
+                # In the edge case (process crash before del lock file is removed),
+                # we do not remove the directory at all.
+                # Since it's in /tmp, this is not that big of a deal.
+                # check if any lock files are remaining
+                remaining_locks = _list_existing_del_locks(temp_dir)
+                if not remaining_locks:
+                    try:
+                        # Timeout 0 means there will be only one attempt to acquire
+                        # the file lock. If it cannot be acquired, a TimeoutError
+                        # will be thrown.
+                        with TempFileLock(temp_dir, timeout=0):
+                            shutil.rmtree(temp_dir, ignore_errors=True)
+                    except TimeoutError:
+                        pass
+
+    def _get_temporary_checkpoint_dir(self) -> str:
+        """Return the name for the temporary checkpoint dir that this checkpoint
+        will get downloaded to, if accessing via `to_directory` or `as_directory`.
+        """
+        tmp_dir_path = tempfile.gettempdir()
+        checkpoint_dir_name = _CHECKPOINT_TEMP_DIR_PREFIX + self._uuid.hex
+        if platform.system() == "Windows":
+            # Max path on Windows is 260 chars, -1 for joining \
+            # Also leave a little for the del lock
+            del_lock_name = _get_del_lock_path("")
+            checkpoint_dir_name = (
+                _CHECKPOINT_TEMP_DIR_PREFIX
+                + self._uuid.hex[
+                    -259
+                    + len(_CHECKPOINT_TEMP_DIR_PREFIX)
+                    + len(tmp_dir_path)
+                    + len(del_lock_name) :
+                ]
+            )
+            if not checkpoint_dir_name.startswith(_CHECKPOINT_TEMP_DIR_PREFIX):
+                raise RuntimeError(
+                    "Couldn't create checkpoint directory due to length "
+                    "constraints. Try specifying a shorter checkpoint path."
+                )
+        return Path(tmp_dir_path, checkpoint_dir_name).as_posix()
+
+    def __fspath__(self):
+        raise TypeError(
+            "You cannot use `Checkpoint` objects directly as paths. "
+            "Use `Checkpoint.to_directory()` or `Checkpoint.as_directory()` instead."
+        )
+
+
+def _get_del_lock_path(path: str, suffix: str = None) -> str:
+    """Get the path to the deletion lock file for a file/directory at `path`.
+
+    Args:
+        path: The path of the file or directory to generate a lock path for.
+        suffix: Suffix appended after ``.del_lock_``. Defaults to the current
+            process ID.
+
+    Returns:
+        The deletion lock file path.
+
+    Example:
+
+        >>> _get_del_lock_path("/tmp/checkpoint_tmp")  # doctest: +ELLIPSIS
+        '/tmp/checkpoint_tmp.del_lock_...
+        >>> _get_del_lock_path("/tmp/checkpoint_tmp/")  # doctest: +ELLIPSIS
+        '/tmp/checkpoint_tmp.del_lock_...
+        >>> _get_del_lock_path("/tmp/checkpoint_tmp.txt")  # doctest: +ELLIPSIS
+        '/tmp/checkpoint_tmp.txt.del_lock_...
+
+    """
+    suffix = suffix if suffix is not None else str(os.getpid())
+    return f"{path.rstrip('/')}.del_lock_{suffix}"
+
+
+def _list_existing_del_locks(path: str) -> List[str]:
+    """List all the deletion lock files for a file/directory at `path`.
+
+    For example, if 2 checkpoints are being read via `as_directory`,
+    then this should return a list of 2 deletion lock files.
+    """
+    return list(glob.glob(f"{_get_del_lock_path(path, suffix='*')}"))
+
+
+def _get_migration_error(name: str):
+    return AttributeError(
+        f"The new `ray.train.Checkpoint` class does not support `{name}()`. "
+        f"Instead, only directories are supported.\n\n"
+        f"Example to store a dictionary in a checkpoint:\n\n"
+        f"import os, tempfile\n"
+        f"import ray.cloudpickle as pickle\n"
+        f"from ray import train\n"
+        f"from ray.train import Checkpoint\n\n"
+        f"with tempfile.TemporaryDirectory() as checkpoint_dir:\n"
+        f"  with open(os.path.join(checkpoint_dir, 'data.pkl'), 'wb') as fp:\n"
+        f"    pickle.dump({{'data': 'value'}}, fp)\n\n"
+        f"  checkpoint = Checkpoint.from_directory(checkpoint_dir)\n"
+        f"  train.report(..., checkpoint=checkpoint)\n\n"
+        f"Example to load a dictionary from a checkpoint:\n\n"
+        f"if train.get_checkpoint():\n"
+        f"  with train.get_checkpoint().as_directory() as checkpoint_dir:\n"
+        f"    with open(os.path.join(checkpoint_dir, 'data.pkl'), 'rb') as fp:\n"
+        f"      data = pickle.load(fp)"
+    )
+
+
+def _get_uri_error(name: str):
+    return AttributeError(
+        f"The new `ray.train.Checkpoint` class does not support `{name}()`. "
+        f"To create a checkpoint from remote storage, create a `Checkpoint` using its "
+        f"constructor instead of `from_directory`.\n"
+        f'Example: `Checkpoint(path="s3://a/b/c")`.\n'
+        f"Then, access the contents of the checkpoint with "
+        f"`checkpoint.as_directory()` / `checkpoint.to_directory()`.\n"
+        f"To upload data to remote storage, use e.g. `pyarrow.fs.FileSystem` "
+        f"or your client of choice."
+    )
+
+
+def _get_preprocessor_error(name: str):
+    return AttributeError(
+        f"The new `ray.train.Checkpoint` class does not support `{name}()`. "
+        f"To include preprocessor information in checkpoints, "
+        f"pass it as metadata in the <Framework>Trainer constructor.\n"
+        f"Example: `TorchTrainer(..., metadata={{...}})`.\n"
+        f"After training, access it in the checkpoint via `checkpoint.get_metadata()`. "
+        f"See here: https://docs.ray.io/en/master/train/user-guides/"
+        f"data-loading-preprocessing.html#preprocessing-structured-data"
+    )

python/ray/train/_internal/accelerator.py

@@ -0,0 +1,5 @@
+import abc
+
+
+class Accelerator(abc.ABC):
+    """A utility that contains methods to accelerate training."""

python/ray/train/_internal/backend_executor.py

@@ -0,0 +1,858 @@
+import logging
+import os
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+)
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._private.accelerators.amd_gpu import HIP_VISIBLE_DEVICES_ENV_VAR
+from ray._private.accelerators.neuron import NEURON_RT_VISIBLE_CORES_ENV_VAR
+from ray._private.accelerators.npu import ASCEND_RT_VISIBLE_DEVICES_ENV_VAR
+from ray._private.accelerators.nvidia_gpu import CUDA_VISIBLE_DEVICES_ENV_VAR
+from ray._private.ray_constants import env_integer
+from ray.exceptions import RayActorError
+from ray.train._checkpoint import Checkpoint
+from ray.train._internal.data_config import DataConfig
+from ray.train._internal.session import (
+    TrialInfo,
+    _TrainingResult,
+    get_session,
+    init_session,
+    shutdown_session,
+)
+from ray.train._internal.storage import StorageContext
+from ray.train._internal.utils import check_for_failure
+from ray.train._internal.worker_group import WorkerGroup
+from ray.train.backend import BackendConfig
+from ray.train.constants import (
+    ENABLE_DETAILED_AUTOFILLED_METRICS_ENV,
+    ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
+    ENABLE_SHARE_HIP_VISIBLE_DEVICES_ENV,
+    ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
+    ENABLE_SHARE_NPU_RT_VISIBLE_DEVICES_ENV,
+    RAY_TRAIN_ENABLE_STATE_TRACKING,
+    TRAIN_ENABLE_WORKER_SPREAD_ENV,
+    TRAIN_PLACEMENT_GROUP_TIMEOUT_S_ENV,
+)
+from ray.util.placement_group import get_current_placement_group, remove_placement_group
+
+if TYPE_CHECKING:
+    from ray.data import Dataset
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+class TrainBackendError(Exception):
+    """Errors with BackendExecutor that should not be exposed to user."""
+
+
+class TrainingWorkerError(Exception):
+    """Raised if a worker fails during training."""
+
+
+@dataclass
+class ResourceConfig:
+    """
+    Resource configuration for resource_ids to share between workers.
+
+    Args:
+        resource_name: The name of the resource to configure
+         (Example: "neuron_cores" or "gpu").
+        resource_enable_sharing_env_var: The environment variable to
+         check if the resource should be shared.
+        share_resource_ids_env_var: The environment variable to configure for
+         sharing the resources with other workers.
+    """
+
+    resource_name: str
+    resource_enable_sharing_env_var: str
+    share_resource_ids_env_var: str
+
+
+class BackendExecutor:
+    """Main execution class for training backends.
+
+    This class holds a worker group and is responsible for executing the
+    training function on the workers, and collecting intermediate results
+    from ``session.report()``.
+
+    Args:
+        backend_config: The configurations for this
+            specific backend.
+        trial_info: Information about the current Tune trial, if running under Tune.
+        num_workers: Number of workers to use for training.
+        resources_per_worker: Dictionary specifying the resources that will be
+            requested for each worker. Defaults to {"CPU": 1}.
+        max_retries: Number of retries when Ray actors fail.
+            Defaults to 3. Set to -1 for unlimited retries.
+    """
+
+    def __init__(
+        self,
+        backend_config: BackendConfig,
+        # TODO(xwjiang): Legacy Ray Train trainer clean up!
+        trial_info: Optional[TrialInfo] = None,
+        num_workers: int = 1,
+        resources_per_worker: Optional[Dict[str, float]] = None,
+        max_retries: int = 3,
+    ):
+        if resources_per_worker is None:
+            self._resources_per_worker = {"CPU": 1}
+        else:
+            self._resources_per_worker = resources_per_worker.copy()
+
+        self._backend_config = backend_config
+        self._backend = backend_config.backend_cls()
+        self._num_workers = num_workers
+        self._max_failures = max_retries
+        if self._max_failures < 0:
+            self._max_failures = float("inf")
+        self._num_failures = 0
+        self._last_failure = None
+        self._initialization_hook = None
+        self._placement_group = None
+
+        self._trial_info = trial_info
+
+        self.worker_group = InactiveWorkerGroup()
+        self.dataset_shards = None
+
+        self._resource_configs = [
+            ResourceConfig(
+                ray_constants.NEURON_CORES,
+                ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
+                NEURON_RT_VISIBLE_CORES_ENV_VAR,
+            ),
+            ResourceConfig(
+                ray_constants.NPU,
+                ENABLE_SHARE_NPU_RT_VISIBLE_DEVICES_ENV,
+                ASCEND_RT_VISIBLE_DEVICES_ENV_VAR,
+            ),
+            # For AMD GPUs, they are using HIP_VISIBLE_DEVICES env var.
+            ResourceConfig(
+                ray_constants.GPU,
+                ENABLE_SHARE_HIP_VISIBLE_DEVICES_ENV,
+                HIP_VISIBLE_DEVICES_ENV_VAR,
+            ),
+        ]
+
+        # Record the initialization time of BackendExecutor, which is
+        # after trainer.fit() and before worker_group executes the training function.
+        self._start_time_ms = int(time.time() * 1000)
+
+        self.state_tracking_enabled = env_integer(RAY_TRAIN_ENABLE_STATE_TRACKING, 0)
+
+    def start(
+        self,
+        initialization_hook: Optional[Callable[[], None]] = None,
+        train_cls: Optional[Type] = None,
+        train_cls_args: Optional[Tuple] = None,
+        train_cls_kwargs: Optional[Dict] = None,
+    ):
+        """Starts the worker group."""
+        self._create_placement_group()
+        placement_group = self._placement_group or "default"
+        self.worker_group = WorkerGroup(
+            num_workers=self._num_workers,
+            resources_per_worker=self._resources_per_worker,
+            actor_cls=train_cls,
+            actor_cls_args=train_cls_args,
+            actor_cls_kwargs=train_cls_kwargs,
+            placement_group=placement_group,
+        )
+        # Hack to avoid OOMs.
+        # This is just a temporary solution for Train loading entire checkpoints
+        # into memory by ensuring that the rank 0 worker is on the same node as
+        # trainable, thus allowing for lazy checkpoint transfer to be used.
+        # See https://github.com/ray-project/ray/issues/33073
+        # for more context.
+        # TODO remove passing in trial_driver_ip.
+
+        trial_driver_node_id = (
+            self._trial_info.driver_node_id if self._trial_info else None
+        )
+        self.worker_group.sort_workers_by_node_id_and_gpu_id(trial_driver_node_id)
+
+        try:
+            if initialization_hook:
+                self._initialization_hook = initialization_hook
+                self.worker_group.execute(initialization_hook)
+
+            # Always propagate the driver's DataContext to each worker in the group.
+            from ray.data import DataContext
+
+            def _set_driver_dataset_context(ctx: DataContext):
+                DataContext._set_current(ctx)
+
+            self.worker_group.execute(
+                _set_driver_dataset_context,
+                DataContext.get_current(),
+            )
+
+            share_cuda_visible_devices_enabled = bool(
+                env_integer(
+                    ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
+                    self._backend.share_cuda_visible_devices,
+                )
+            )
+
+            if (
+                self._resources_per_worker.get("GPU", 0) > 0
+                and share_cuda_visible_devices_enabled
+            ):
+                self._share_cuda_visible_devices()
+            for resource_config in self._resource_configs:
+                if self._is_share_resources_enabled(
+                    resource_config.resource_name,
+                    resource_config.resource_enable_sharing_env_var,
+                ):
+                    self._share_resource_ids(
+                        resource_config.resource_name,
+                        resource_config.share_resource_ids_env_var,
+                    )
+            self._backend.on_start(self.worker_group, self._backend_config)
+        except RayActorError as exc:
+            logger.exception(str(exc))
+            logger.warning(
+                "Failure occurred during startup. Restarting all workers and "
+                "attempting to startup again."
+            )
+            self._increment_failures()
+            self._restart()
+
+        if self.state_tracking_enabled:
+            from ray.train._internal.state import TrainRunStateManager
+            from ray.train._internal.state.state_actor import get_state_actor
+
+            self.state_manager = TrainRunStateManager(state_actor=get_state_actor())
+
+    def _create_placement_group(self):
+        """Creates a placement group if it does not exist.
+
+        If a placement group is already detected (Tune) this will be a no-op.
+
+        By default the placement group will be created with PACK strategy.
+        This is optimized for colocating GPUs on a minimal number of nodes.
+        This behavior can be overridden to use the SPREAD strategy by defining
+        ``TRAIN_ENABLE_WORKER_SPREAD_ENV``
+
+        If a placement group is created it will be stored as
+        self._placement_group.
+        """
+        current_placement_group = get_current_placement_group()
+        worker = ray._private.worker.global_worker
+        should_capture_child_tasks_in_placement_group = (
+            worker.should_capture_child_tasks_in_placement_group
+        )
+        should_create_placement_group = (
+            current_placement_group is None
+            or not should_capture_child_tasks_in_placement_group
+        )
+
+        if should_create_placement_group:
+            bundles = [
+                self._resources_per_worker.copy() for _ in range(self._num_workers)
+            ]
+
+            use_spread = bool(env_integer(TRAIN_ENABLE_WORKER_SPREAD_ENV, 0))
+            strategy = "SPREAD" if use_spread else "PACK"
+
+            placement_group = ray.util.placement_group(bundles, strategy=strategy)
+            logger.debug("Waiting for placement group to start.")
+            timeout = env_integer(TRAIN_PLACEMENT_GROUP_TIMEOUT_S_ENV, 100)
+            ready, _ = ray.wait([placement_group.ready()], timeout=timeout)
+            if ready:
+                logger.debug("Placement group has started.")
+            else:
+                raise TimeoutError(
+                    "Placement group creation timed out. Make sure your "
+                    "cluster either has enough resources or use an "
+                    "autoscaling cluster. If you are running on a cluster, "
+                    "make sure you specify an address in `ray.init()`, for example, "
+                    '`ray.init("auto")`. You can also increase the timeout by setting '
+                    "the TRAIN_PLACEMENT_GROUP_TIMEOUT_S environment variable. "
+                    "Current resources available: {}, resources requested by the "
+                    "placement group: {}".format(
+                        ray.available_resources(), placement_group.bundle_specs
+                    )
+                )
+            self._placement_group = placement_group
+
+    def _share_cuda_visible_devices(self):
+        """Sets CUDA_VISIBLE_DEVICES on all workers.
+
+        For each worker, CUDA_VISIBLE_DEVICES will be set to the GPU IDs
+        visible to all workers on that worker's node.
+
+        This allows GPU workers on the same node to communicate with one
+        another.
+
+        Example:
+
+            Setup:
+            - Node1:
+                - Worker1: {0, 1}
+                - Worker2: {2, 3}
+            - Node2:
+                - Worker3: {0, 1}
+
+            CUDA_VISIBLE_DEVICES:
+            - Worker1: "0,1,2,3"
+            - Worker2: "0,1,2,3"
+            - Worker3: "0,1"
+
+        """
+        self._share_resource_ids(ray_constants.GPU, CUDA_VISIBLE_DEVICES_ENV_VAR)
+
+    def _share_resource_ids(self, resource: str, env_var: str):
+        """Sets the given env_var on all workers.
+
+        For each worker, the cores/devices are visible to all the
+        workers on that worker's node.This allows workers on the
+        same node to communicate with one another.
+
+        Example:
+
+            Setup:
+            - Node1:
+                - Worker1: {0, 1}
+                - Worker2: {2, 3}
+            - Node2:
+                - Worker3: {0, 1}
+
+            NEURON_RT_VISIBLE_CORES/TPU_VISIBLE_CHIPS/...:
+            - Worker1: "0,1,2,3"
+            - Worker2: "0,1,2,3"
+            - Worker2: "0,1"
+
+        Args:
+            resource: The name of the resource/accelerator.
+            env_var: The name of the environment variable to set.
+        """
+        node_ids_and_resource_ids = [
+            (
+                w.metadata.node_id,
+                w.metadata.resource_ids[resource],
+            )
+            for w in self.worker_group.workers
+        ]
+        node_id_to_worker_id = defaultdict(set)
+        node_id_to_resource_ids = defaultdict(set)
+
+        for worker_id, (node_id, resource_ids) in enumerate(node_ids_and_resource_ids):
+            node_id_to_worker_id[node_id].add(worker_id)
+            node_id_to_resource_ids[node_id].update(resource_ids)
+
+        futures = []
+        for node_id, resource_ids in node_id_to_resource_ids.items():
+            resource_ids = sorted(resource_ids)
+            all_resource_ids = ",".join(resource_ids)
+
+            def set_resource_ids():
+                os.environ[env_var] = all_resource_ids
+
+            for worker_id in node_id_to_worker_id[node_id]:
+                futures.append(
+                    self.worker_group.execute_single_async(worker_id, set_resource_ids)
+                )
+        ray.get(futures)
+
+    def _is_share_resources_enabled(self, resource_name: str, enable_sharing_env: str):
+        """Whether to share resource IDs on all workers
+        based on enable_sharing_env.
+
+        This will return true if resources are requested and greater than 0.
+        Also, user can disable by configuring the `enable_sharing_env` to "0".
+
+        Args:
+            resource_name: The name of the resource/accelerator.
+            enable_sharing_env: The name of the environment variable
+                to check.
+
+        Returns:
+            True if resource sharing is enabled, False otherwise.
+        """
+        has_resource_requested = self._resources_per_worker.get(resource_name, 0) > 0
+        return has_resource_requested and ray_constants.env_bool(
+            enable_sharing_env, True
+        )
+
+    def _create_rank_world_size_mappings(
+        self,
+    ) -> Tuple[Dict[int, int], Dict[int, int], Dict[int, int]]:
+        """Create rank and world size mappings for workers.
+
+        There are three maps returned:
+            - local_rank_map, which maps from worker world_rank to local_rank.
+            - local_world_size_map, which maps from world_rank to local_world_size
+            - node_rank_map, which maps from world rank to node rank
+
+        Example:
+            Worker 0: node 0
+            Worker 1: node 0
+            Worker 2: node 1
+            Worker 3: node 0
+            Worker 4: node 1
+
+            Workers 0, 1, 3 are on node 0.
+            Workers 2, 4 are on node 1.
+
+            Expected local_rank_map:
+            {
+                0 -> 0,
+                1 -> 1,
+                2 -> 0,
+                3 -> 2,
+                4 -> 1
+            }
+
+            Expected local_world_size_map:
+            {
+                0 -> 3,
+                1 -> 3,
+                2 -> 2,
+                3 -> 3,
+                4 -> 2
+            }
+
+            Expected node_rank_map:
+            {
+                0 -> 0,
+                1 -> 0,
+                2 -> 1,
+                3 -> 0,
+                4 -> 1
+            }
+
+        Returns:
+            A tuple of (local_rank_map, local_world_size_map, node_rank_map).
+        """
+        local_rank_map = {}  # map from world rank to local rank
+        local_world_size_map = {}  # map from world rank to local world size
+        node_rank_map = {}  # map from world rank to node rank
+        node_ids = {}  # map from node id to node index
+        node_cnt = 0  # count the number of nodes
+
+        node_id_dict = defaultdict(
+            int
+        )  # map from node id to the number of workers on it.
+        for world_rank in range(len(self.worker_group)):
+            worker = self.worker_group.workers[world_rank]
+            node_id = worker.metadata.node_id
+            local_rank_map[world_rank] = node_id_dict[node_id]
+            node_id_dict[node_id] += 1
+
+            if node_id not in node_ids:
+                node_ids[node_id] = node_cnt
+                node_cnt += 1
+            node_rank_map[world_rank] = node_ids[node_id]
+
+        for world_rank in range(len(self.worker_group)):
+            worker = self.worker_group.workers[world_rank]
+            node_id = worker.metadata.node_id
+            local_world_size_map[world_rank] = node_id_dict[node_id]
+
+        workers_info = "\n".join(
+            [
+                f"- (node_id={w.metadata.node_id}, ip={w.metadata.node_ip}, "
+                f"pid={w.metadata.pid}) world_rank={i}, "
+                f"local_rank={local_rank_map[i]}, node_rank={node_rank_map[i]}"
+                for i, w in enumerate(self.worker_group.workers)
+            ]
+        )
+        logger.info(f"Started distributed worker processes: \n{workers_info}")
+
+        return local_rank_map, local_world_size_map, node_rank_map
+
+    def start_training(
+        self,
+        train_func: Callable[[], T],
+        datasets: Dict[str, "Dataset"],
+        metadata: Dict[str, Any],
+        data_config: DataConfig,
+        storage: StorageContext,
+        checkpoint: Optional[Checkpoint] = None,
+    ) -> None:
+        """Executes a training function on all workers in a separate thread.
+
+        ``finish_training`` should be called after this.
+
+        Args:
+            train_func: The training function to run on each worker.
+            datasets: The base datasets.
+            metadata: User-supplied metadata dict propagated to checkpoints
+                created during training.
+            data_config: The config object for creating dataset shards for workers.
+            storage: The storage context, providing access to the experiment
+                directory and other persistent storage state.
+            checkpoint: The checkpoint data that
+                should be loaded onto each worker and accessed by the
+                training function via ``session.get_checkpoint()``. If this
+                is ``None`` then no checkpoint will be loaded.
+        """
+        use_detailed_autofilled_metrics = env_integer(
+            ENABLE_DETAILED_AUTOFILLED_METRICS_ENV, 0
+        )
+
+        # First initialize the session.
+        def initialize_session(
+            train_func,
+            world_rank,
+            local_rank,
+            node_rank,
+            local_world_size,
+            world_size,
+            trial_info,
+            checkpoint,
+            dataset_shard,
+            metadata,
+            storage,
+        ):
+            try:
+                init_session(
+                    training_func=train_func,
+                    world_rank=world_rank,
+                    local_rank=local_rank,
+                    node_rank=node_rank,
+                    local_world_size=local_world_size,
+                    world_size=world_size,
+                    trial_info=trial_info,
+                    dataset_shard=dataset_shard,
+                    metadata=metadata,
+                    checkpoint=checkpoint,
+                    detailed_autofilled_metrics=use_detailed_autofilled_metrics,
+                    storage=storage,
+                )
+            except ValueError:
+                raise TrainBackendError(
+                    "Attempting to start training but a "
+                    "previous training run is still ongoing. "
+                    "You must call `finish_training` before "
+                    "calling `start_training` again."
+                )
+
+        if self.dataset_shards is None:
+            actors = [worker.actor for worker in self.worker_group.workers]
+            node_ids = [worker.metadata.node_id for worker in self.worker_group.workers]
+            self.dataset_shards = data_config.configure(
+                datasets,
+                world_size=len(self.worker_group),
+                worker_handles=actors,
+                worker_node_ids=node_ids,
+            )
+
+        (
+            local_rank_map,
+            local_world_size_map,
+            node_rank_map,
+        ) = self._create_rank_world_size_mappings()
+
+        futures = []
+        for index in range(len(self.worker_group)):
+            futures.append(
+                self.worker_group.execute_single_async(
+                    index,
+                    initialize_session,
+                    world_rank=index,
+                    local_rank=local_rank_map[index],
+                    node_rank=node_rank_map[index],
+                    local_world_size=local_world_size_map[index],
+                    world_size=len(self.worker_group),
+                    trial_info=self._trial_info,
+                    train_func=train_func,
+                    dataset_shard=self.dataset_shards[index],
+                    metadata=metadata,
+                    checkpoint=checkpoint,
+                    storage=storage,
+                )
+            )
+
+        self._backend.on_training_start(self.worker_group, self._backend_config)
+
+        self.get_with_failure_handling(futures)
+
+        # Register Train Run before training starts
+        if self.state_tracking_enabled:
+            from ray.train._internal.state.schema import RunStatusEnum
+
+            core_context = ray.runtime_context.get_runtime_context()
+
+            self.state_manager.register_train_run(
+                run_id=self._trial_info.run_id,
+                run_name=self._trial_info.experiment_name,
+                job_id=core_context.get_job_id(),
+                controller_actor_id=core_context.get_actor_id(),
+                datasets=datasets,
+                worker_group=self.worker_group,
+                start_time_ms=self._start_time_ms,
+                run_status=RunStatusEnum.RUNNING,
+                resources=[self._resources_per_worker] * self._num_workers,
+            )
+
+        # Run the training function asynchronously in its own thread.
+        def train_async():
+            session = get_session()
+            session.start()
+
+        self.worker_group.execute_async(train_async)
+
+    def get_next_results(self) -> Optional[List[_TrainingResult]]:
+        """Fetches the next ``_TrainingResult`` from each worker.
+
+        Each ``_TrainingResult`` is expected to correspond to the same step from
+        each worker (e.g. the same call to ``train.report()``).
+
+        Returns:
+            A list of ``_TrainingResult``s or ``None`` if there are no more results
+            since the training function has exited on all workers.
+        """
+
+        def get_next():
+            session = _get_session("get_next_results")
+            try:
+                result = session.get_next()
+            except RuntimeError:
+                # Training thread has not been started yet.
+                raise TrainBackendError(
+                    "`get_next_results` has been called "
+                    "before `start_training`. Please call "
+                    "`start_training` before "
+                    "`get_next_results`."
+                )
+
+            return result
+
+        # Get next result from each worker.
+        futures = self.worker_group.execute_async(get_next)
+        results = self.get_with_failure_handling(futures)
+
+        # Check if any worker returned None.
+        if any(r is None for r in results):
+            # Either all workers have results or none of them do.
+            if not all(r is None for r in results):
+                raise RuntimeError(
+                    "Some workers returned results while "
+                    "others didn't. Make sure that "
+                    "`session.report()` are called the "
+                    "same number of times on all workers."
+                )
+            else:
+                # Return None if all results are None.
+                return None
+
+        return results
+
+    def pause_reporting(self):
+        """Disable workers from enqueuing results from ``session.report()``.
+
+        Note: Already reported results may still be enqueued at this point,
+              and should be handled appropriately.
+        """
+
+        def pause_session_reporting():
+            session = _get_session("pause_reporting")
+            return session.pause_reporting()
+
+        futures = self.worker_group.execute_async(pause_session_reporting)
+        self.get_with_failure_handling(futures)
+
+    def finish_training(self):
+        """Finish training and return final results. Propagate any exceptions.
+
+        Blocks until training is finished on all workers.
+
+        Assumes `start_training` has already been called.
+
+        Returns:
+            A list of return values from calling ``train_func`` on each worker.
+                Each item corresponds to the return value from a single worker.
+        """
+
+        def end_training():
+            session = _get_session("finish_training")
+            try:
+                # session.finish raises any Exceptions from training.
+                output = session.finish()
+            finally:
+                # Shutdown session even if session.finish() raises an
+                # Exception.
+                shutdown_session()
+
+            return output
+
+        futures = self.worker_group.execute_async(end_training)
+        results = self.get_with_failure_handling(futures)
+        return results
+
+    def report_final_run_status(
+        self,
+        errored: bool = False,
+        failed_rank: Optional[int] = None,
+        stack_trace: Optional[str] = None,
+    ):
+        """Report the final train run status, error, and end time to TrainStateActor."""
+        if self.state_tracking_enabled:
+            from ray.train._internal.state.schema import (
+                MAX_ERROR_STACK_TRACE_LENGTH,
+                RunStatusEnum,
+            )
+
+            if errored:
+                run_status = RunStatusEnum.ERRORED
+                status_detail = ""
+                if failed_rank is not None:
+                    status_detail += f"Rank {failed_rank} worker raised an error. \n"
+                if stack_trace is not None:
+                    # Keep only the last part of the stack trace if it's too long.
+                    status_detail += stack_trace[-MAX_ERROR_STACK_TRACE_LENGTH:]
+            else:
+                run_status = RunStatusEnum.FINISHED
+                status_detail = ""
+
+            self.state_manager.end_train_run(
+                run_id=self._trial_info.run_id,
+                run_status=run_status,
+                status_detail=status_detail,
+                end_time_ms=int(time.time() * 1000),
+            )
+
+    def get_with_failure_handling(self, remote_values: List[ray.ObjectRef]):
+        """Gets the remote values while handling for worker failures.
+
+        This method should be called instead of ``ray.get()`` directly in
+        order to handle worker failures.
+
+        If a worker failure is identified, backend specific failure handling
+        is executed and a ``TrainingWorkerError`` is raised.
+
+        Args:
+            remote_values: List of object refs representing functions
+                that may fail in the middle of execution. For example, running
+                a Train training loop in multiple parallel actor calls.
+        Returns:
+            The resolved objects represented by the passed in ObjectRefs.
+        """
+        success, exception = check_for_failure(remote_values)
+        if success:
+            return ray.get(remote_values)
+        else:
+            self._last_failure = exception
+            self._increment_failures()
+            logger.warning(
+                "Failure identified during training. Restarting all workers and "
+                "continuing training from latest checkpoint."
+            )
+            self._restart()
+            raise TrainingWorkerError
+
+    def shutdown(self, graceful_termination: bool = True):
+        """Shuts down the workers in the worker group.
+
+        Args:
+            graceful_termination: If set to True, attempt to clean up the backend
+                before terminating the Ray actors.
+
+        """
+        if graceful_termination:
+            try:
+                self._backend.on_shutdown(self.worker_group, self._backend_config)
+            except RayActorError:
+                logger.warning(
+                    "Graceful shutdown of backend failed. This is "
+                    "expected if one of the workers has crashed."
+                )
+
+        if graceful_termination:
+            self.worker_group.shutdown()
+        else:
+            self.worker_group.shutdown(patience_s=0)
+        self.worker_group = InactiveWorkerGroup()
+
+        if self._placement_group:
+            remove_placement_group(self._placement_group)
+            self._placement_group = None
+
+        self.dataset_shards = None
+
+    def is_started(self):
+        return not isinstance(self.worker_group, InactiveWorkerGroup)
+
+    def _restart(self):
+        self.worker_group.shutdown()
+        if self._initialization_hook is not None:
+            initialization_hook = self._initialization_hook
+        else:
+            initialization_hook = None
+        if self._placement_group:
+            remove_placement_group(self._placement_group)
+            self._placement_group = None
+        self.start(initialization_hook=initialization_hook)
+
+    def _increment_failures(self):
+        self._num_failures += 1
+        if self._num_failures >= self._max_failures:
+            failure = self._last_failure
+            self._last_failure = None
+            if self._max_failures > 0:
+                exc = RuntimeError(
+                    f"Training has failed after {self._num_failures} attempts."
+                )
+                raise exc.with_traceback(None) from failure
+            else:
+                raise failure
+
+    def get_worker_group(self):
+        return self.worker_group
+
+    def _get_num_failures(self):
+        return self._num_failures
+
+
+class InactiveWorkerGroupError(Exception):
+    """Raised when underlying worker group is inactive."""
+
+
+class InactiveWorkerGroup:
+    # TODO: fix inheritence. perhaps create WorkerGroupInterface.
+
+    # Need to define getstate and setstate so that getattr does not screwup
+    # pickling. See https://stackoverflow.com/a/50888571/11249691
+    def __getstate__(self):
+        return vars(self)
+
+    def __setstate__(self, state):
+        vars(self).update(state)
+
+    def __getattr__(self, name):
+        raise InactiveWorkerGroupError()
+
+    def __len__(self):
+        raise InactiveWorkerGroupError()
+
+
+def _get_session(method_name: str):
+    # Get the session for this worker.
+    session = get_session()
+    if not session:
+        # Session is not initialized yet.
+        raise TrainBackendError(
+            f"`{method_name}` has been called "
+            "before `start_training`. Please call "
+            "`start_training` before "
+            f"`{method_name}`."
+        )
+    return session

python/ray/train/_internal/base_worker_group.py

@@ -0,0 +1,105 @@
+"""Abstract base class for WorkerGroup implementations.
+
+This module defines the common base class that both V1 and V2 WorkerGroup
+implementations should inherit from to ensure backend compatibility.
+"""
+
+import abc
+from typing import Callable, List, TypeVar
+
+from ray.types import ObjectRef
+from ray.util.annotations import DeveloperAPI
+
+T = TypeVar("T")
+
+
+@DeveloperAPI
+class BaseWorkerGroup(abc.ABC):
+    """Abstract base class for WorkerGroup implementations.
+
+    This base class defines the minimal set of methods that backend classes
+    expect from WorkerGroup implementations. Both V1 and V2 WorkerGroup
+    classes should inherit from this base class to ensure compatibility with
+    all backend configurations.
+
+    The interface focuses on the core operations that backends need:
+    - Executing functions on workers
+    - Getting worker count and resource allocation
+    """
+
+    @abc.abstractmethod
+    def execute(self, func: Callable[..., T], *args, **kwargs) -> List[T]:
+        """Execute a function on all workers synchronously.
+
+        Args:
+            func: The function to execute on each worker.
+            *args: Positional arguments to pass to the function.
+            **kwargs: Keyword arguments to pass to the function.
+
+        Returns:
+            A list of results from each worker, in worker rank order.
+        """
+        pass
+
+    @abc.abstractmethod
+    def execute_async(self, func: Callable[..., T], *args, **kwargs) -> List[ObjectRef]:
+        """Execute a function on all workers asynchronously.
+
+        Args:
+            func: The function to execute on each worker.
+            *args: Positional arguments to pass to the function.
+            **kwargs: Keyword arguments to pass to the function.
+
+        Returns:
+            A list of ObjectRef results from each worker, in worker rank order.
+        """
+        pass
+
+    @abc.abstractmethod
+    def execute_single(
+        self, worker_index: int, func: Callable[..., T], *args, **kwargs
+    ) -> T:
+        """Execute a function on a single worker synchronously.
+
+        Args:
+            worker_index: The index of the worker to execute on.
+            func: The function to execute.
+            *args: Positional arguments to pass to the function.
+            **kwargs: Keyword arguments to pass to the function.
+
+        Returns:
+            The result from the specified worker.
+        """
+        pass
+
+    @abc.abstractmethod
+    def execute_single_async(
+        self, worker_index: int, func: Callable[..., T], *args, **kwargs
+    ) -> ObjectRef:
+        """Execute a function on a single worker asynchronously.
+
+        Args:
+            worker_index: The index of the worker to execute on.
+            func: The function to execute.
+            *args: Positional arguments to pass to the function.
+            **kwargs: Keyword arguments to pass to the function.
+
+        Returns:
+            An ObjectRef to the result from the specified worker.
+        """
+        pass
+
+    @abc.abstractmethod
+    def __len__(self) -> int:
+        """Return the number of workers in the group."""
+        pass
+
+    @abc.abstractmethod
+    def get_resources_per_worker(self) -> dict:
+        """Get the resources allocated per worker.
+
+        Returns:
+            A dictionary mapping resource names to quantities per worker.
+            Common keys include "CPU", "GPU", "memory".
+        """
+        pass

python/ray/train/_internal/checkpoint_manager.py

@@ -0,0 +1,206 @@
+import logging
+import numbers
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from ray._private import ray_constants
+from ray._private.dict import flatten_dict
+from ray.air._internal.util import is_nan
+from ray.air.config import MAX
+from ray.train import Checkpoint, CheckpointConfig
+from ray.train._internal.session import _TrainingResult
+from ray.train._internal.storage import _delete_fs_path
+from ray.train.constants import TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE
+
+logger = logging.getLogger(__name__)
+
+
+def _insert_into_sorted_list(
+    list: List[_TrainingResult],
+    item: _TrainingResult,
+    key: Callable[[_TrainingResult], Any],
+    checkpoint_to_report_index: Optional[Dict[Checkpoint, int]] = None,
+):
+    """Insert an item into a sorted list with a custom key function.
+
+    Args:
+        list: The list to insert the item into.
+        item: The item to insert.
+        key: The key function to use to sort the list.
+        checkpoint_to_report_index: A dictionary mapping checkpoints to report indices.
+            Used to break ties when scores are equal.
+    """
+    checkpoint_to_report_index = checkpoint_to_report_index or {}
+    # TODO: optimize this with sortedlist, batching, etc
+    i = 0
+    while i < len(list):
+        # When scores are equal, later checkpoints are later in the list.
+        list_item_key, item_key = key(list[i]), key(item)
+        if list_item_key > item_key or (
+            list_item_key == item_key
+            and checkpoint_to_report_index.get(list[i].checkpoint, 0)
+            > checkpoint_to_report_index.get(item.checkpoint, 0)
+        ):
+            break
+        i += 1
+    list.insert(i, item)
+
+
+class _CheckpointManager:
+    """Checkpoint manager that handles checkpoint book-keeping for a trial.
+
+    The main purpose of this abstraction is to keep the top K checkpoints based on
+    recency/a user-provided metric.
+
+    NOTE: This class interacts with `_TrainingResult` objects, which are
+    (checkpoint, metrics) pairs. This is to order checkpoints by metrics.
+
+    Args:
+        checkpoint_config: Defines how many and which checkpoints to keep.
+    """
+
+    def __init__(self, checkpoint_config: Optional[CheckpointConfig]):
+        self._checkpoint_config = checkpoint_config or CheckpointConfig()
+
+        # List of checkpoints ordered by ascending score.
+        self._checkpoint_results: List[_TrainingResult] = []
+
+        # The latest registered checkpoint.
+        # This should never be immediately deleted upon registration,
+        # even if it's not in the top K checkpoints, based on score.
+        self._latest_checkpoint_result: Optional[_TrainingResult] = None
+
+        if (
+            self._checkpoint_config.num_to_keep is not None
+            and self._checkpoint_config.num_to_keep <= 0
+        ):
+            raise ValueError(
+                f"`num_to_keep` must >= 1, got: "
+                f"{self._checkpoint_config.num_to_keep}"
+            )
+
+    @property
+    def checkpoint_config(self):
+        return self._checkpoint_config
+
+    def register_checkpoint(self, checkpoint_result: _TrainingResult):
+        """Register new checkpoint and add to bookkeeping.
+
+        This method will register a new checkpoint and add it to the internal
+        bookkeeping logic. This means the checkpoint manager will decide if
+        this checkpoint should be kept, and if older or worse performing
+        checkpoints should be deleted.
+
+        Args:
+            checkpoint_result: Tracked training result containing the checkpoint
+                and associated metrics to add to bookkeeping.
+        """
+        self._latest_checkpoint_result = checkpoint_result
+
+        score_attr = self._checkpoint_config.checkpoint_score_attribute
+        if ray_constants.env_bool(TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE, False):
+            metrics = (
+                {score_attr: checkpoint_result.metrics[score_attr]}
+                if score_attr in checkpoint_result.metrics
+                else {}
+            )
+            checkpoint_result = _TrainingResult(
+                checkpoint=checkpoint_result.checkpoint,
+                metrics=metrics,
+            )
+
+        if score_attr is not None and score_attr in checkpoint_result.metrics:
+            # If we're ordering by a score, insert the checkpoint
+            # so that the list remains sorted.
+            _insert_into_sorted_list(
+                self._checkpoint_results,
+                checkpoint_result,
+                key=self._get_checkpoint_score,
+            )
+        else:
+            # If no metric is provided, just append (ordering by time of registration).
+            self._checkpoint_results.append(checkpoint_result)
+
+        if self._checkpoint_config.num_to_keep is not None:
+            # Delete the bottom (N - K) checkpoints
+            worst_results = set(
+                self._checkpoint_results[: -self._checkpoint_config.num_to_keep]
+            )
+            # Except for the latest checkpoint.
+            results_to_delete = worst_results - {self._latest_checkpoint_result}
+
+            # Update internal state before actually deleting them.
+            self._checkpoint_results = [
+                checkpoint_result
+                for checkpoint_result in self._checkpoint_results
+                if checkpoint_result not in results_to_delete
+            ]
+
+            for checkpoint_result in results_to_delete:
+                checkpoint = checkpoint_result.checkpoint
+                logger.debug("Deleting checkpoint: %s", checkpoint)
+                _delete_fs_path(fs=checkpoint.filesystem, fs_path=checkpoint.path)
+
+    def _get_checkpoint_score(
+        self, checkpoint: _TrainingResult
+    ) -> Tuple[bool, numbers.Number]:
+        """Get the score for a checkpoint, according to checkpoint config.
+
+        If `mode="min"`, the metric is negated so that the lowest score is
+        treated as the best.
+
+        Args:
+            checkpoint: The training result whose metrics should be scored.
+
+        Returns:
+            Tuple: A tuple of (not_is_nan: bool, score: numbers.Number).
+                This score orders: nan values < float("-inf") < valid numeric metrics
+        """
+        checkpoint_score_attribute = self._checkpoint_config.checkpoint_score_attribute
+        if checkpoint_score_attribute:
+            flat_metrics = flatten_dict(checkpoint.metrics)
+            try:
+                checkpoint_result = flat_metrics[checkpoint_score_attribute]
+            except KeyError:
+                valid_keys = list(flat_metrics.keys())
+                logger.error(
+                    f"Result dict has no key: {checkpoint_score_attribute}. "
+                    f"checkpoint_score_attr must be set to a key in the "
+                    f"result dict. Valid keys are: {valid_keys}"
+                )
+                checkpoint_result = float("-inf")
+        else:
+            checkpoint_result = float("-inf")
+
+        checkpoint_score_order = self._checkpoint_config.checkpoint_score_order
+        order_factor = 1.0 if checkpoint_score_order == MAX else -1.0
+
+        checkpoint_score = order_factor * checkpoint_result
+
+        if not isinstance(checkpoint_score, numbers.Number):
+            raise ValueError(
+                f"Unable to persist checkpoint for "
+                f"checkpoint_score_attribute: "
+                f"{checkpoint_score_attribute} with value "
+                f"{checkpoint_score}. "
+                f"This attribute must be numerical."
+            )
+
+        return (
+            (not is_nan(checkpoint_score), checkpoint_score)
+            if not is_nan(checkpoint_score)
+            else (False, float("-inf"))
+        )
+
+    @property
+    def best_checkpoint_result(self) -> Optional[_TrainingResult]:
+        return self._checkpoint_results[-1] if self._checkpoint_results else None
+
+    @property
+    def latest_checkpoint_result(self) -> Optional[_TrainingResult]:
+        return self._latest_checkpoint_result
+
+    @property
+    def best_checkpoint_results(self) -> List[_TrainingResult]:
+        if self._checkpoint_config.num_to_keep is None:
+            return self._checkpoint_results
+        return self._checkpoint_results[-self._checkpoint_config.num_to_keep :]

python/ray/train/_internal/data_config.py

@@ -0,0 +1,178 @@
+import copy
+from collections import defaultdict
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
+
+from ray.actor import ActorHandle
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data import DataIterator, Dataset, ExecutionOptions, NodeIdStr
+
+
+@PublicAPI(stability="stable")
+class DataConfig:
+    """Class responsible for configuring Train dataset preprocessing.
+
+    For advanced use cases, this class can be subclassed and the `configure()` method
+    overridden for custom data preprocessing.
+    """
+
+    def __init__(
+        self,
+        datasets_to_split: Union[Literal["all"], List[str]] = "all",
+        execution_options: Optional[
+            Union["ExecutionOptions", Dict[str, "ExecutionOptions"]]
+        ] = None,
+        enable_shard_locality: bool = True,
+    ):
+        """Construct a DataConfig.
+
+        Args:
+            datasets_to_split: Specifies which datasets should be split among workers.
+                Can be set to "all" or a list of dataset names. Defaults to "all",
+                i.e. split all datasets.
+            execution_options: The execution options to pass to Ray Data. Can be either:
+                1. A single ExecutionOptions object that is applied to all datasets.
+                2. A dict mapping dataset names to ExecutionOptions. If a dataset name
+                is not in the dict, it defaults to ``DataConfig.default_ingest_options()``.
+                By default, the options are optimized for data ingest. When overriding,
+                base your options off ``DataConfig.default_ingest_options()``.
+            enable_shard_locality: If true, dataset sharding across Train workers will
+                consider locality to minimize cross-node data transfer. Enabled by default.
+        """
+        from ray.data import ExecutionOptions
+
+        if isinstance(datasets_to_split, list) or datasets_to_split == "all":
+            self._datasets_to_split = datasets_to_split
+        else:
+            raise TypeError(
+                "`datasets_to_split` should be a 'all' or a list of strings of "
+                "dataset names. Received "
+                f"{type(datasets_to_split).__name__} with value {datasets_to_split}."
+            )
+
+        default_execution_options = DataConfig.default_ingest_options()
+
+        if isinstance(execution_options, ExecutionOptions):
+            default_execution_options = execution_options
+        # If None, all datasets will use the default ingest options.
+        self._execution_options: Dict[str, "ExecutionOptions"] = defaultdict(
+            lambda: copy.deepcopy(default_execution_options)
+        )
+        if isinstance(execution_options, dict):
+            self._execution_options.update(execution_options)
+
+        self._enable_shard_locality = enable_shard_locality
+
+        self._num_train_cpus = 0.0
+        self._num_train_gpus = 0.0
+
+    def set_train_total_resources(self, num_train_cpus: float, num_train_gpus: float):
+        """Set the total number of CPUs and GPUs used by training.
+
+        If CPU or GPU resource limits are not set, they will be set to the
+        total cluster resources minus the resources used by training.
+        """
+        # TODO: We may also include other resources besides CPU and GPU.
+        self._num_train_cpus = num_train_cpus
+        self._num_train_gpus = num_train_gpus
+
+    def _get_execution_options(self, dataset_name: str) -> "ExecutionOptions":
+        """Return a copy of the configured execution options for a given dataset name."""
+        return copy.deepcopy(self._execution_options[dataset_name])
+
+    @DeveloperAPI
+    def configure(
+        self,
+        datasets: Dict[str, "Dataset"],
+        world_size: int,
+        worker_handles: Optional[List[ActorHandle]],
+        worker_node_ids: Optional[List["NodeIdStr"]],
+        **kwargs,
+    ) -> List[Dict[str, "DataIterator"]]:
+        """Configure how Train datasets should be assigned to workers.
+
+        Args:
+            datasets: The datasets dict passed to Train by the user.
+            world_size: The number of Train workers in total.
+            worker_handles: The actor handles of the Train workers.
+            worker_node_ids: The node ids of the Train workers.
+            **kwargs: Forwards compatibility placeholder.
+
+        Returns:
+            A list of dataset splits for each worker. The size of the list must be
+            equal to `world_size`. Each element of the list contains the assigned
+            `DataIterator` instances by name for the worker.
+        """
+        from ray.data._internal.execution.interfaces.execution_options import (
+            ExecutionResources,
+        )
+
+        output = [{} for _ in range(world_size)]
+
+        for dataset_name, dataset in datasets.items():
+            if dataset.name is None:
+                dataset.set_name(dataset_name)
+
+        if self._datasets_to_split == "all":
+            datasets_to_split = set(datasets.keys())
+        else:
+            datasets_to_split = set(self._datasets_to_split)
+
+        locality_hints = worker_node_ids if self._enable_shard_locality else None
+        for name, ds in datasets.items():
+            execution_options = self._get_execution_options(name)
+
+            if execution_options.is_resource_limits_default():
+                if not self._scaling_policy_reserves_train_resources():
+                    execution_options.exclude_resources = (
+                        execution_options.exclude_resources.add(
+                            ExecutionResources(
+                                cpu=self._num_train_cpus, gpu=self._num_train_gpus
+                            )
+                        )
+                    )
+
+            ds = ds.copy(ds)
+            ds.context.execution_options = execution_options
+
+            if name in datasets_to_split:
+                for i, split in enumerate(
+                    ds.streaming_split(
+                        world_size, equal=True, locality_hints=locality_hints
+                    )
+                ):
+                    output[i][name] = split
+            else:
+                for i in range(world_size):
+                    output[i][name] = ds.iterator()
+
+        return output
+
+    @classmethod
+    def _scaling_policy_reserves_train_resources(cls) -> bool:
+        """True iff Ray Train V2's ScalingPolicy will register training resources
+        with the AutoscalingCoordinator for this run.
+
+        """
+        from ray.train.v2._internal.constants import is_v2_enabled
+
+        return is_v2_enabled()
+
+    @staticmethod
+    def default_ingest_options() -> "ExecutionOptions":
+        """The default Ray Data options used for data ingest.
+
+        By default, configurations are carried over from what is already set
+        in DataContext.
+        """
+        from ray.data import ExecutionOptions
+        from ray.data.context import DataContext
+
+        ctx = DataContext.get_current()
+        return ExecutionOptions(
+            resource_limits=ctx.execution_options.resource_limits,
+            exclude_resources=ctx.execution_options.exclude_resources,
+            preserve_order=ctx.execution_options.preserve_order,
+            verbose_progress=ctx.execution_options.verbose_progress,
+        )

python/ray/train/_internal/framework_checkpoint.py

@@ -0,0 +1,57 @@
+from typing import TYPE_CHECKING, Optional
+
+import ray.cloudpickle as ray_pickle
+from ray._common.utils import binary_to_hex, hex_to_binary
+from ray.train._checkpoint import Checkpoint
+
+if TYPE_CHECKING:
+    from ray.data.preprocessor import Preprocessor
+
+PREPROCESSOR_KEY = "preprocessor_pkl"
+
+
+class FrameworkCheckpoint(Checkpoint):
+    """A checkpoint to preserve the functionality of legacy
+    framework-specific checkpoints.
+
+    Example:
+
+        >>> import tempfile
+        >>> from ray.data.preprocessor import Preprocessor
+        >>> checkpoint = FrameworkCheckpoint(tempfile.mkdtemp())
+        >>> checkpoint.get_preprocessor() is None
+        True
+        >>> preprocessor = Preprocessor()
+        >>> preprocessor._attr = 1234
+        >>> checkpoint.set_preprocessor(preprocessor)
+        >>> checkpoint.get_preprocessor()._attr
+        1234
+    """
+
+    def get_preprocessor(self) -> Optional["Preprocessor"]:
+        """Return the preprocessor stored in the checkpoint.
+
+        .. warning::
+
+            The checkpoint path must point to a **trusted** source.
+            The preprocessor is stored as a pickle blob inside the checkpoint
+            metadata. Loading a checkpoint from an untrusted path (shared
+            storage, downloaded artifact, checkpoint produced by a different
+            party) is equivalent to executing arbitrary Python code. Never
+            call this method on a checkpoint you do not fully control.
+
+        Returns:
+            The preprocessor stored in the checkpoint, or ``None`` if no
+            preprocessor was stored.
+        """
+        metadata = self.get_metadata()
+        preprocessor_bytes = metadata.get(PREPROCESSOR_KEY)
+        if preprocessor_bytes is None:
+            return None
+        return ray_pickle.loads(hex_to_binary(preprocessor_bytes))
+
+    def set_preprocessor(self, preprocessor: "Preprocessor"):
+        """Store a preprocessor with the checkpoint."""
+        self.update_metadata(
+            {PREPROCESSOR_KEY: binary_to_hex(ray_pickle.dumps(preprocessor))}
+        )

python/ray/train/_internal/state/__init__.py

@@ -0,0 +1,14 @@
+from ray.train._internal.state.state_manager import TrainRunStateManager
+
+try:
+    import pydantic  # noqa: F401
+except ImportError:
+    raise ModuleNotFoundError(
+        "pydantic isn't installed."
+        "To install pydantic, please run 'pip install pydantic'"
+    )
+
+
+__all__ = [
+    "TrainRunStateManager",
+]

python/ray/train/_internal/state/export.py

@@ -0,0 +1,105 @@
+from typing import Optional
+
+from ray.core.generated.export_train_state_pb2 import (
+    ExportTrainRunAttemptEventData as ProtoTrainRunAttempt,
+    ExportTrainRunEventData as ProtoTrainRun,
+)
+from ray.train._internal.state.schema import (
+    ActorStatusEnum,
+    RunStatusEnum,
+    TrainRunInfo,
+    TrainWorkerInfo,
+)
+
+TRAIN_SCHEMA_VERSION = 1
+RAY_TRAIN_VERSION = 1
+
+# Status mapping dictionaries
+_ACTOR_STATUS_MAP = {
+    ActorStatusEnum.ALIVE: ProtoTrainRunAttempt.ActorStatus.ALIVE,
+    ActorStatusEnum.DEAD: ProtoTrainRunAttempt.ActorStatus.DEAD,
+}
+
+_RUN_ATTEMPT_STATUS_MAP = {
+    RunStatusEnum.STARTED: ProtoTrainRunAttempt.RunAttemptStatus.PENDING,
+    RunStatusEnum.RUNNING: ProtoTrainRunAttempt.RunAttemptStatus.RUNNING,
+    RunStatusEnum.FINISHED: ProtoTrainRunAttempt.RunAttemptStatus.FINISHED,
+    RunStatusEnum.ERRORED: ProtoTrainRunAttempt.RunAttemptStatus.ERRORED,
+    RunStatusEnum.ABORTED: ProtoTrainRunAttempt.RunAttemptStatus.ABORTED,
+}
+
+_RUN_STATUS_MAP = {
+    RunStatusEnum.STARTED: ProtoTrainRun.RunStatus.INITIALIZING,
+    RunStatusEnum.RUNNING: ProtoTrainRun.RunStatus.RUNNING,
+    RunStatusEnum.FINISHED: ProtoTrainRun.RunStatus.FINISHED,
+    RunStatusEnum.ERRORED: ProtoTrainRun.RunStatus.ERRORED,
+    RunStatusEnum.ABORTED: ProtoTrainRun.RunStatus.ABORTED,
+}
+
+
+def _ms_to_ns(ms: Optional[int]) -> Optional[int]:
+    if ms is None:
+        return None
+    return ms * 1000000
+
+
+# Helper conversion functions
+def _to_proto_resources(resources: dict) -> ProtoTrainRunAttempt.TrainResources:
+    """Convert resources dictionary to protobuf TrainResources."""
+    return ProtoTrainRunAttempt.TrainResources(resources=resources)
+
+
+def _to_proto_worker(worker: TrainWorkerInfo) -> ProtoTrainRunAttempt.TrainWorker:
+    """Convert TrainWorker to protobuf format."""
+    proto_worker = ProtoTrainRunAttempt.TrainWorker(
+        world_rank=worker.world_rank,
+        local_rank=worker.local_rank,
+        node_rank=worker.node_rank,
+        actor_id=bytes.fromhex(worker.actor_id),
+        node_id=bytes.fromhex(worker.node_id),
+        node_ip=worker.node_ip,
+        pid=worker.pid,
+        gpu_ids=worker.gpu_ids,
+        status=_ACTOR_STATUS_MAP[worker.status],
+        resources=_to_proto_resources(worker.resources),
+    )
+
+    return proto_worker
+
+
+# Main conversion functions
+def train_run_info_to_proto_run(run_info: TrainRunInfo) -> ProtoTrainRun:
+    """Convert TrainRunInfo to TrainRun protobuf format."""
+    proto_run = ProtoTrainRun(
+        schema_version=TRAIN_SCHEMA_VERSION,
+        ray_train_version=RAY_TRAIN_VERSION,
+        id=run_info.id,
+        name=run_info.name,
+        job_id=bytes.fromhex(run_info.job_id),
+        controller_actor_id=bytes.fromhex(run_info.controller_actor_id),
+        status=_RUN_STATUS_MAP[run_info.run_status],
+        status_detail=run_info.status_detail,
+        start_time_ns=_ms_to_ns(run_info.start_time_ms),
+        end_time_ns=_ms_to_ns(run_info.end_time_ms),
+    )
+
+    return proto_run
+
+
+def train_run_info_to_proto_attempt(run_info: TrainRunInfo) -> ProtoTrainRunAttempt:
+    """Convert TrainRunInfo to TrainRunAttempt protobuf format."""
+
+    proto_attempt = ProtoTrainRunAttempt(
+        schema_version=TRAIN_SCHEMA_VERSION,
+        ray_train_version=RAY_TRAIN_VERSION,
+        run_id=run_info.id,
+        attempt_id=run_info.id,  # Same as run_id
+        status=_RUN_ATTEMPT_STATUS_MAP[run_info.run_status],
+        status_detail=run_info.status_detail,
+        start_time_ns=_ms_to_ns(run_info.start_time_ms),
+        end_time_ns=_ms_to_ns(run_info.end_time_ms),
+        resources=[_to_proto_resources(r) for r in run_info.resources],
+        workers=[_to_proto_worker(worker) for worker in run_info.workers],
+    )
+
+    return proto_attempt

python/ray/train/_internal/state/schema.py

@@ -0,0 +1,165 @@
+from enum import Enum
+from typing import Dict, List, Optional
+
+from ray._common.pydantic_compat import BaseModel, Field
+from ray.dashboard.modules.job.pydantic_models import JobDetails
+from ray.util.annotations import DeveloperAPI
+
+MAX_ERROR_STACK_TRACE_LENGTH = 50000
+
+
+@DeveloperAPI
+class RunStatusEnum(str, Enum):
+    """Enumeration for the status of a train run."""
+
+    # (Deprecated) Replaced by RUNNING.
+    # The train run has started
+    STARTED = "STARTED"
+    # The train run is running
+    RUNNING = "RUNNING"
+    # The train run was terminated as expected
+    FINISHED = "FINISHED"
+    # The train run was terminated early due to errors in the training function
+    ERRORED = "ERRORED"
+    # The train run was terminated early due to system errors or controller errors
+    ABORTED = "ABORTED"
+
+
+@DeveloperAPI
+class ActorStatusEnum(str, Enum):
+    DEAD = "DEAD"
+    ALIVE = "ALIVE"
+
+
+@DeveloperAPI
+class TrainWorkerInfo(BaseModel):
+    """Metadata of a Ray Train worker."""
+
+    actor_id: str = Field(description="Actor ID of the worker.")
+    world_rank: int = Field(description="World rank of the worker.")
+    local_rank: int = Field(description="Local rank of the worker.")
+    node_rank: int = Field(description="Node rank of the worker.")
+    node_id: str = Field(description="ID of the node that the worker is running on.")
+    node_ip: str = Field(
+        description="IP address of the node that the worker is running on."
+    )
+    pid: int = Field(description="Process ID of the worker.")
+    gpu_ids: List[int] = Field(
+        description="A list of GPU ids allocated to that worker."
+    )
+    status: ActorStatusEnum = Field(
+        description="The status of the train worker actor. It can be ALIVE or DEAD."
+    )
+    resources: Dict[str, float] = Field(
+        description="The resources allocated to the worker."
+    )
+
+
+@DeveloperAPI
+class MemoryInfo(BaseModel):
+    rss: int
+    vms: int
+    pfaults: Optional[int] = None
+    pageins: Optional[int] = None
+
+
+@DeveloperAPI
+class ProcessStats(BaseModel):
+    cpuPercent: float
+    # total memory, free memory, memory used ratio
+    mem: Optional[List[int]] = None
+    memoryInfo: MemoryInfo
+
+
+class ProcessGPUUsage(BaseModel):
+    # This gpu usage stats from a process
+    pid: int
+    gpuMemoryUsage: int
+
+
+@DeveloperAPI
+class GPUStats(BaseModel):
+    uuid: str
+    index: int
+    name: str
+    utilizationGpu: Optional[float] = None
+    memoryUsed: float
+    memoryTotal: float
+    processInfo: ProcessGPUUsage
+
+
+@DeveloperAPI
+class TrainWorkerInfoWithDetails(TrainWorkerInfo):
+    """Metadata of a Ray Train worker."""
+
+    processStats: Optional[ProcessStats] = Field(
+        None, description="Process stats of the worker."
+    )
+    gpus: List[GPUStats] = Field(
+        default_factory=list,
+        description=(
+            "GPU stats of the worker. "
+            "Only returns GPUs that are attached to the worker process."
+        ),
+    )
+
+
+@DeveloperAPI
+class TrainDatasetInfo(BaseModel):
+    name: str = Field(
+        description="The key of the dataset dict specified in Ray Train Trainer."
+    )
+    dataset_uuid: str = Field(description="The uuid of the dataset.")
+    dataset_name: Optional[str] = Field(None, description="The name of the dataset.")
+
+
+@DeveloperAPI
+class TrainRunInfo(BaseModel):
+    """Metadata for a Ray Train run and information about its workers."""
+
+    name: str = Field(description="The name of the Train run.")
+    id: str = Field(description="The unique identifier for each Train run.")
+    job_id: str = Field(description="The Ray Job ID.")
+    controller_actor_id: str = Field(description="Actor Id of the Train controller.")
+    workers: List[TrainWorkerInfo] = Field(
+        description="A List of Train workers sorted by global ranks."
+    )
+    datasets: List[TrainDatasetInfo] = Field(
+        description="A List of dataset info for this Train run."
+    )
+    run_status: RunStatusEnum = Field(
+        description="The current status of the train run. It can be one of the "
+        "following: RUNNING, FINISHED, ERRORED, or ABORTED."
+    )
+    status_detail: str = Field(
+        description="Detailed information about the current run status, "
+        "such as error messages."
+    )
+    start_time_ms: int = Field(
+        description="The UNIX timestamp of the start time of this Train run."
+    )
+    end_time_ms: Optional[int] = Field(
+        None,
+        description="The UNIX timestamp of the end time of this Train run. "
+        "If null, the Train run has not ended yet.",
+    )
+    resources: List[Dict[str, float]] = Field(
+        description="The resources allocated to the worker."
+    )
+
+
+@DeveloperAPI
+class TrainRunInfoWithDetails(TrainRunInfo):
+    """Metadata for a Ray Train run and information about its workers."""
+
+    workers: List[TrainWorkerInfoWithDetails] = Field(
+        description="A List of Train workers sorted by global ranks."
+    )
+    job_details: Optional[JobDetails] = Field(
+        None, description="Details of the job that started this Train run."
+    )
+
+
+@DeveloperAPI
+class TrainRunsResponse(BaseModel):
+    train_runs: List[TrainRunInfoWithDetails]

python/ray/train/_internal/state/state_actor.py

@@ -0,0 +1,152 @@
+import logging
+import os
+import threading
+from typing import Dict, Optional
+
+import ray
+from ray._private.event.export_event_logger import (
+    EventLogType,
+    check_export_api_enabled,
+    get_export_event_logger,
+)
+from ray.actor import ActorHandle
+from ray.train._internal.state.schema import TrainRunInfo
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote(num_cpus=0)
+class TrainStateActor:
+    def __init__(self):
+        self._run_infos: Dict[str, TrainRunInfo] = {}
+
+        (
+            self._export_logger,
+            self._is_train_run_export_api_enabled,
+            self._is_train_run_attempt_export_api_enabled,
+        ) = self._init_export_logger()
+
+    def register_train_run(self, run_info: TrainRunInfo) -> None:
+        # Register a new train run.
+        self._run_infos[run_info.id] = run_info
+
+        self._maybe_export_train_run(run_info)
+        self._maybe_export_train_run_attempt(run_info)
+
+    def get_train_run(self, run_id: str) -> Optional[TrainRunInfo]:
+        # Retrieve a registered run with its id
+        return self._run_infos.get(run_id, None)
+
+    def get_all_train_runs(self) -> Dict[str, TrainRunInfo]:
+        # Retrieve all registered train runs
+        return self._run_infos
+
+    # ============================
+    # Export API
+    # ============================
+
+    def is_export_api_enabled(self) -> bool:
+        return self._export_logger is not None
+
+    def _init_export_logger(self) -> tuple[Optional[logging.Logger], bool, bool]:
+        """Initialize the export logger and check if the export API is enabled.
+
+        Returns:
+            A tuple containing:
+                - The export logger (or None if export API is not enabled).
+                - A boolean indicating if the export API is enabled for train runs.
+                - A boolean indicating if the export API is enabled for train run attempts.
+        """
+        # Proto schemas should be imported within the scope of TrainStateActor to
+        # prevent serialization errors.
+        from ray.core.generated.export_event_pb2 import ExportEvent
+
+        is_train_run_export_api_enabled = check_export_api_enabled(
+            ExportEvent.SourceType.EXPORT_TRAIN_RUN
+        )
+        is_train_run_attempt_export_api_enabled = check_export_api_enabled(
+            ExportEvent.SourceType.EXPORT_TRAIN_RUN_ATTEMPT
+        )
+        export_api_enabled = (
+            is_train_run_export_api_enabled or is_train_run_attempt_export_api_enabled
+        )
+
+        if not export_api_enabled:
+            return None, False, False
+
+        log_directory = os.path.join(
+            ray._private.worker._global_node.get_session_dir_path(), "logs"
+        )
+        logger = None
+        try:
+            logger = get_export_event_logger(
+                EventLogType.TRAIN_STATE,
+                log_directory,
+            )
+        except Exception:
+            logger.exception(
+                "Unable to initialize the export event logger, so no Train export "
+                "events will be written."
+            )
+
+        if logger is None:
+            return None, False, False
+
+        return (
+            logger,
+            is_train_run_export_api_enabled,
+            is_train_run_attempt_export_api_enabled,
+        )
+
+    def _maybe_export_train_run(self, run_info: TrainRunInfo) -> None:
+        if not self._is_train_run_export_api_enabled:
+            return
+
+        from ray.train._internal.state.export import train_run_info_to_proto_run
+
+        run_proto = train_run_info_to_proto_run(run_info)
+        self._export_logger.send_event(run_proto)
+
+    def _maybe_export_train_run_attempt(self, run_info: TrainRunInfo) -> None:
+        if not self._is_train_run_attempt_export_api_enabled:
+            return
+
+        from ray.train._internal.state.export import train_run_info_to_proto_attempt
+
+        run_attempt_proto = train_run_info_to_proto_attempt(run_info)
+        self._export_logger.send_event(run_attempt_proto)
+
+
+TRAIN_STATE_ACTOR_NAME = "train_state_actor"
+TRAIN_STATE_ACTOR_NAMESPACE = "_train_state_actor"
+
+_state_actor_lock: threading.RLock = threading.RLock()
+
+
+def get_or_create_state_actor() -> ActorHandle:
+    """Get or create a `TrainStateActor` on the head node."""
+    with _state_actor_lock:
+        state_actor = TrainStateActor.options(
+            name=TRAIN_STATE_ACTOR_NAME,
+            namespace=TRAIN_STATE_ACTOR_NAMESPACE,
+            get_if_exists=True,
+            lifetime="detached",
+            resources={"node:__internal_head__": 0.001},
+            # Escape from the parent's placement group
+            scheduling_strategy="DEFAULT",
+        ).remote()
+
+    # Ensure the state actor is ready
+    ray.get(state_actor.__ray_ready__.remote())
+    return state_actor
+
+
+def get_state_actor() -> Optional[ActorHandle]:
+    """Get the `TrainStateActor` if exists, otherwise return None."""
+    try:
+        return ray.get_actor(
+            name=TRAIN_STATE_ACTOR_NAME,
+            namespace=TRAIN_STATE_ACTOR_NAMESPACE,
+        )
+    except ValueError:
+        return None

python/ray/train/_internal/state/state_manager.py

@@ -0,0 +1,132 @@
+import logging
+import os
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Dict, List
+
+import ray
+from ray.train._internal.state.schema import (
+    ActorStatusEnum,
+    RunStatusEnum,
+    TrainDatasetInfo,
+    TrainRunInfo,
+    TrainWorkerInfo,
+)
+from ray.train._internal.utils import check_for_failure
+from ray.train._internal.worker_group import WorkerGroup
+
+if TYPE_CHECKING:
+    from ray.data import Dataset
+
+logger = logging.getLogger(__name__)
+
+
+class TrainRunStateManager:
+    """A class that aggregates and reports train run info to TrainStateActor.
+
+    This manager class is created on the train controller layer for each run.
+    """
+
+    def __init__(self, state_actor) -> None:
+        self.state_actor = state_actor
+        self.train_run_info_dict = defaultdict(dict)
+
+    def register_train_run(
+        self,
+        run_id: str,
+        job_id: str,
+        run_name: str,
+        run_status: str,
+        controller_actor_id: str,
+        datasets: Dict[str, "Dataset"],
+        worker_group: WorkerGroup,
+        start_time_ms: float,
+        resources: List[Dict[str, float]],
+        status_detail: str = "",
+    ) -> None:
+        """Collect Train Run Info and report to StateActor."""
+
+        if not self.state_actor:
+            logger.warning(
+                "Unable to register train run since `TrainStateActor` is not started."
+            )
+            return
+
+        def collect_train_worker_info():
+            train_context = ray.train.get_context()
+            core_context = ray.runtime_context.get_runtime_context()
+            return TrainWorkerInfo(
+                world_rank=train_context.get_world_rank(),
+                local_rank=train_context.get_local_rank(),
+                node_rank=train_context.get_node_rank(),
+                actor_id=core_context.get_actor_id(),
+                node_id=core_context.get_node_id(),
+                node_ip=ray.util.get_node_ip_address(),
+                gpu_ids=ray.get_gpu_ids(),
+                pid=os.getpid(),
+                resources=resources[0],
+                status=ActorStatusEnum.ALIVE,
+            )
+
+        futures = [
+            worker_group.execute_single_async(index, collect_train_worker_info)
+            for index in range(len(worker_group))
+        ]
+        success, exception = check_for_failure(futures)
+
+        if not success:
+            logger.error(
+                "Failed to collect run information from the Ray Train "
+                f"workers:\n{exception}"
+            )
+            return
+
+        worker_info_list = ray.get(futures)
+        worker_info_list = sorted(worker_info_list, key=lambda info: info.world_rank)
+
+        dataset_info_list = [
+            TrainDatasetInfo(
+                name=ds_name,
+                dataset_name=ds._dataset_name,
+                dataset_uuid=ds._uuid,
+            )
+            for ds_name, ds in datasets.items()
+        ]
+
+        updates = dict(
+            id=run_id,
+            job_id=job_id,
+            name=run_name,
+            controller_actor_id=controller_actor_id,
+            workers=worker_info_list,
+            datasets=dataset_info_list,
+            start_time_ms=start_time_ms,
+            run_status=run_status,
+            status_detail=status_detail,
+            resources=resources,
+        )
+
+        # Clear the cached info to avoid registering the same run twice
+        self.train_run_info_dict[run_id] = {}
+        self._update_train_run_info(run_id, updates)
+
+    def end_train_run(
+        self,
+        run_id: str,
+        run_status: RunStatusEnum,
+        status_detail: str,
+        end_time_ms: int,
+    ):
+        """Update the train run status when the training is finished."""
+        updates = dict(
+            run_status=run_status,
+            status_detail=status_detail,
+            end_time_ms=end_time_ms,
+        )
+        self._update_train_run_info(run_id, updates)
+
+    def _update_train_run_info(self, run_id: str, updates: Dict[str, Any]) -> None:
+        """Update specific fields of a registered TrainRunInfo instance."""
+        if run_id in self.train_run_info_dict:
+            self.train_run_info_dict[run_id].update(updates)
+            train_run_info = TrainRunInfo(**self.train_run_info_dict[run_id])
+            ray.get(self.state_actor.register_train_run.remote(train_run_info))

python/ray/train/_internal/storage.py

@@ -0,0 +1,736 @@
+# Try import ray[train] core requirements (defined in setup.py)
+# isort: off
+try:
+    import fsspec  # noqa
+    from fsspec.implementations.local import LocalFileSystem
+
+except (ImportError, ModuleNotFoundError) as e:
+    raise RuntimeError(
+        "fsspec is a required dependency of Ray Train and Ray Tune. "
+        "Please install with: `pip install fsspec`"
+    ) from e
+
+try:
+    import pyarrow
+    import pyarrow.fs
+
+except (ImportError, ModuleNotFoundError) as e:
+    raise RuntimeError(
+        "pyarrow is a required dependency of Ray Train and Ray Tune. "
+        "Please install with: `pip install pyarrow`"
+    ) from e
+
+try:
+    # check if Arrow has S3 support
+    from pyarrow.fs import S3FileSystem
+except ImportError:
+    S3FileSystem = None
+# isort: on
+
+import fnmatch
+import logging
+import os
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Type, Union
+
+from ray.air._internal.filelock import TempFileLock
+from ray.train._internal.syncer import SyncConfig, Syncer, _BackgroundSyncer
+from ray.train.constants import _get_ray_train_session_dir
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.train._checkpoint import Checkpoint
+
+
+logger = logging.getLogger(__name__)
+
+
+_VALIDATE_STORAGE_MARKER_FILENAME = ".validate_storage_marker"
+
+
+class _ExcludingLocalFilesystem(LocalFileSystem):
+    """LocalFileSystem wrapper to exclude files according to patterns.
+
+    Args:
+        root_path: Root path to strip when matching with the exclude pattern.
+            Ex: root_path="/tmp/a/b/c", exclude=["*a*"], will exclude
+            /tmp/a/b/c/_a_.txt but not ALL of /tmp/a/*.
+        exclude: List of patterns that are applied to files returned by
+            ``self.find()``. If a file path matches this pattern, it will
+            be excluded.
+        **kwargs: Forwarded to the ``fsspec.implementations.local.LocalFileSystem``
+            parent class.
+
+    """
+
+    def __init__(self, root_path: Path, exclude: List[str], **kwargs):
+        super().__init__(**kwargs)
+        self._exclude = exclude
+        self._root_path = root_path
+
+    @property
+    def fsid(self):
+        return "_excluding_local"
+
+    def _should_exclude(self, path: str) -> bool:
+        """Return True if `path` (relative to `root_path`) matches any of the
+        `self._exclude` patterns."""
+        path = Path(path)
+        relative_path = path.relative_to(self._root_path).as_posix()
+        match_candidates = [relative_path]
+        if path.is_dir():
+            # Everything is in posix path format ('/')
+            match_candidates.append(relative_path + "/")
+
+        for excl in self._exclude:
+            if any(fnmatch.fnmatch(candidate, excl) for candidate in match_candidates):
+                return True
+        return False
+
+    def find(self, path, maxdepth=None, withdirs=False, detail=False, **kwargs):
+        """Call parent find() and exclude from result."""
+        paths = super().find(
+            path, maxdepth=maxdepth, withdirs=withdirs, detail=detail, **kwargs
+        )
+        if detail:
+            return {
+                path: out
+                for path, out in paths.items()
+                if not self._should_exclude(path)
+            }
+        else:
+            return [path for path in paths if not self._should_exclude(path)]
+
+
+def _pyarrow_fs_copy_files(
+    source, destination, source_filesystem=None, destination_filesystem=None, **kwargs
+):
+    if S3FileSystem and isinstance(destination_filesystem, pyarrow.fs.S3FileSystem):
+        # Workaround multi-threading issue with pyarrow. Note that use_threads=True
+        # is safe for download, just not for uploads, see:
+        # https://github.com/apache/arrow/issues/32372
+        kwargs.setdefault("use_threads", False)
+
+    # Use a large chunk size to speed up large checkpoint transfers.
+    kwargs.setdefault("chunk_size", 64 * 1024 * 1024)
+
+    return pyarrow.fs.copy_files(
+        source,
+        destination,
+        source_filesystem=source_filesystem,
+        destination_filesystem=destination_filesystem,
+        **kwargs,
+    )
+
+
+# TODO(justinvyu): Add unit tests for all these utils.
+
+
+def _delete_fs_path(fs: pyarrow.fs.FileSystem, fs_path: str):
+    is_dir = _is_directory(fs, fs_path)
+
+    try:
+        if is_dir:
+            fs.delete_dir(fs_path)
+        else:
+            fs.delete_file(fs_path)
+    except Exception:
+        logger.exception(f"Caught exception when deleting path at ({fs}, {fs_path}):")
+
+
+def _download_from_fs_path(
+    fs: pyarrow.fs.FileSystem,
+    fs_path: str,
+    local_path: str,
+    filelock: bool = True,
+):
+    """Downloads a directory or file from (fs, fs_path) to a local path.
+
+    If fs_path points to a directory:
+    - The full directory contents are downloaded directly into `local_path`,
+      rather than to a subdirectory of `local_path`.
+
+    If fs_path points to a file:
+    - The file is downloaded to `local_path`, which is expected to be a file path.
+
+    If the download fails, the `local_path` contents are
+    cleaned up before raising, if the directory did not previously exist.
+
+    NOTE: This method creates `local_path`'s parent directories if they do not
+    already exist. If the download fails, this does NOT clean up all the parent
+    directories that were created.
+
+    Args:
+        fs: The filesystem to download from.
+        fs_path: The filesystem path (either a directory or a file) to download.
+        local_path: The local path to download to.
+        filelock: Whether to require a file lock before downloading, useful for
+            multiple downloads to the same directory that may be happening in parallel.
+
+    Raises:
+        FileNotFoundError: if (fs, fs_path) doesn't exist.
+    """
+
+    _local_path = Path(local_path).resolve()
+    exists_before = _local_path.exists()
+    if _is_directory(fs=fs, fs_path=fs_path):
+        _local_path.mkdir(parents=True, exist_ok=True)
+    else:
+        _local_path.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        if filelock:
+            with TempFileLock(f"{os.path.normpath(local_path)}.lock"):
+                _pyarrow_fs_copy_files(fs_path, local_path, source_filesystem=fs)
+        else:
+            _pyarrow_fs_copy_files(fs_path, local_path, source_filesystem=fs)
+    except Exception as e:
+        # Clean up the directory if downloading was unsuccessful
+        if not exists_before:
+            shutil.rmtree(local_path, ignore_errors=True)
+        raise e
+
+
+def _upload_to_fs_path(
+    local_path: str,
+    fs: pyarrow.fs.FileSystem,
+    fs_path: str,
+    exclude: Optional[List[str]] = None,
+) -> None:
+    """Uploads a local directory or file to (fs, fs_path).
+
+    NOTE: This will create all necessary parent directories at the destination.
+
+    Args:
+        local_path: The local path to upload.
+        fs: The filesystem to upload to.
+        fs_path: The filesystem path where the dir/file will be uploaded to.
+        exclude: A list of filename matches to exclude from upload. This includes
+            all files under subdirectories as well.
+            This pattern will match with the relative paths of all files under
+            `local_path`.
+            Ex: ["*.png"] to exclude all .png images.
+    """
+
+    if not exclude:
+        # TODO(justinvyu): uploading a single file doesn't work
+        # (since we always create a directory at fs_path)
+        _create_directory(fs=fs, fs_path=fs_path)
+        _pyarrow_fs_copy_files(local_path, fs_path, destination_filesystem=fs)
+        return
+
+    _upload_to_uri_with_exclude_fsspec(
+        local_path=local_path, fs=fs, fs_path=fs_path, exclude=exclude
+    )
+
+
+def _upload_to_uri_with_exclude_fsspec(
+    local_path: str, fs: "pyarrow.fs", fs_path: str, exclude: Optional[List[str]]
+) -> None:
+    local_fs = _ExcludingLocalFilesystem(root_path=local_path, exclude=exclude)
+    handler = pyarrow.fs.FSSpecHandler(local_fs)
+    source_fs = pyarrow.fs.PyFileSystem(handler)
+
+    _create_directory(fs=fs, fs_path=fs_path)
+    _pyarrow_fs_copy_files(
+        local_path, fs_path, source_filesystem=source_fs, destination_filesystem=fs
+    )
+
+
+def _list_at_fs_path(
+    fs: pyarrow.fs.FileSystem,
+    fs_path: str,
+    file_filter: Optional[Callable[[pyarrow.fs.FileInfo], bool]] = None,
+) -> List[str]:
+    """Returns the list of filenames at (fs, fs_path), similar to os.listdir.
+
+    If the path doesn't exist, returns an empty list.
+    """
+    if file_filter is None:
+        file_filter = lambda x: True  # noqa: E731
+
+    selector = pyarrow.fs.FileSelector(fs_path, allow_not_found=True, recursive=False)
+    return [
+        os.path.relpath(file_info.path.lstrip("/"), start=fs_path.lstrip("/"))
+        for file_info in fs.get_file_info(selector)
+        if file_filter(file_info)
+    ]
+
+
+def _exists_at_fs_path(fs: pyarrow.fs.FileSystem, fs_path: str) -> bool:
+    """Returns True if (fs, fs_path) exists."""
+
+    valid = fs.get_file_info(fs_path)
+    return valid.type != pyarrow.fs.FileType.NotFound
+
+
+def _is_directory(fs: pyarrow.fs.FileSystem, fs_path: str) -> bool:
+    """Checks if (fs, fs_path) is a directory or a file.
+
+    Args:
+        fs: The filesystem to use.
+        fs_path: The path on the filesystem to check.
+
+    Returns:
+        True if the path is a directory, False if it is a file.
+
+    Raises:
+        FileNotFoundError: if (fs, fs_path) doesn't exist.
+    """
+
+    file_info = fs.get_file_info(fs_path)
+    if file_info.type == pyarrow.fs.FileType.NotFound:
+        raise FileNotFoundError(f"Path not found: ({fs}, {fs_path})")
+
+    return not file_info.is_file
+
+
+def _create_directory(fs: pyarrow.fs.FileSystem, fs_path: str) -> None:
+    """Create directory at (fs, fs_path).
+
+    Some external filesystems require directories to already exist, or at least
+    the `netloc` to be created (e.g. PyArrows ``mock://`` filesystem).
+
+    Generally this should be done before and outside of Ray applications. This
+    utility is thus primarily used in testing, e.g. of ``mock://` URIs.
+    """
+    try:
+        fs.create_dir(fs_path)
+    except Exception:
+        logger.exception(
+            f"Caught exception when creating directory at ({fs}, {fs_path}):"
+        )
+
+
+def get_fs_and_path(
+    storage_path: Union[str, os.PathLike],
+    storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
+) -> Tuple[pyarrow.fs.FileSystem, str]:
+    """Returns the fs and path from a storage path and an optional custom fs.
+
+    Args:
+        storage_path: A storage path or URI. (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 storage_path
+            is assumed to be prefix-stripped already, and must be a valid path
+            on the filesystem.
+
+    Returns:
+        A tuple of (filesystem, path) resolved from the inputs.
+    """
+    storage_path = str(storage_path)
+
+    if storage_filesystem:
+        return storage_filesystem, storage_path
+
+    return pyarrow.fs.FileSystem.from_uri(storage_path)
+
+
+class _FilesystemSyncer(_BackgroundSyncer):
+    """Syncer between local filesystem and a `storage_filesystem`."""
+
+    def __init__(self, storage_filesystem: Optional["pyarrow.fs.FileSystem"], **kwargs):
+        self.storage_filesystem = storage_filesystem
+        super().__init__(**kwargs)
+
+    def _sync_up_command(
+        self, local_path: str, uri: str, exclude: Optional[List] = None
+    ) -> Tuple[Callable, Dict]:
+        # TODO(justinvyu): Defer this cleanup up as part of the
+        # external-facing Syncer deprecation.
+        fs_path = uri
+        return (
+            _upload_to_fs_path,
+            dict(
+                local_path=local_path,
+                fs=self.storage_filesystem,
+                fs_path=fs_path,
+                exclude=exclude,
+            ),
+        )
+
+    def _sync_down_command(self, uri: str, local_path: str) -> Tuple[Callable, Dict]:
+        fs_path = uri
+        return (
+            _download_from_fs_path,
+            dict(
+                fs=self.storage_filesystem,
+                fs_path=fs_path,
+                local_path=local_path,
+            ),
+        )
+
+    def _delete_command(self, uri: str) -> Tuple[Callable, Dict]:
+        fs_path = uri
+        return _delete_fs_path, dict(fs=self.storage_filesystem, fs_path=fs_path)
+
+
+@DeveloperAPI
+class StorageContext:
+    """Shared context that holds the source of truth for all paths and
+    storage utilities, passed along from the driver to workers.
+
+    This object defines a few types of paths:
+    1. *_fs_path: A path on the `storage_filesystem`. This is a regular path
+        which has been prefix-stripped by pyarrow.fs.FileSystem.from_uri and
+        can be joined with `Path(...).as_posix()`.
+    2. *_driver_staging_path: The temporary staging directory on the local filesystem
+        where driver artifacts are saved to before persisting them to storage.
+    3. trial_working_directory: The local filesystem path that the remote
+        actors' working directories are moved to by default.
+        This is separated from the driver staging path so that driver syncing
+        does not implicitly upload the trial working directory, for trials on the
+        driver node.
+
+    Example with storage_path="mock:///bucket/path?param=1":
+
+        >>> import ray
+        >>> from ray.train._internal.storage import StorageContext
+        >>> import os
+        >>> _ = ray.init()
+        >>> storage = StorageContext(
+        ...     storage_path="mock://netloc/bucket/path?param=1",
+        ...     experiment_dir_name="exp_name",
+        ... )
+        >>> storage.storage_filesystem   # Auto-resolved  # doctest: +ELLIPSIS
+        <pyarrow._fs._MockFileSystem object...
+        >>> storage.experiment_fs_path
+        'bucket/path/exp_name'
+        >>> storage.experiment_driver_staging_path  # doctest: +ELLIPSIS
+        '/tmp/ray/session_.../artifacts/.../exp_name/driver_artifacts'
+        >>> storage.trial_dir_name = "trial_dir"
+        >>> storage.trial_fs_path
+        'bucket/path/exp_name/trial_dir'
+        >>> storage.trial_driver_staging_path  # doctest: +ELLIPSIS
+        '/tmp/ray/session_.../artifacts/.../exp_name/driver_artifacts/trial_dir'
+        >>> storage.trial_working_directory   # doctest: +ELLIPSIS
+        '/tmp/ray/session_.../artifacts/.../exp_name/working_dirs/trial_dir'
+        >>> storage.current_checkpoint_index = 1
+        >>> storage.checkpoint_fs_path
+        'bucket/path/exp_name/trial_dir/checkpoint_000001'
+        >>> ray.shutdown()
+
+    Example with storage_path="/tmp/ray_results":
+
+        >>> from ray.train._internal.storage import StorageContext
+        >>> storage = StorageContext(
+        ...     storage_path="/tmp/ray_results",
+        ...     experiment_dir_name="exp_name",
+        ... )
+        >>> storage.storage_fs_path
+        '/tmp/ray_results'
+        >>> storage.experiment_fs_path
+        '/tmp/ray_results/exp_name'
+        >>> storage.storage_filesystem   # Auto-resolved  # doctest: +ELLIPSIS
+        <pyarrow._fs.LocalFileSystem object...
+
+    Internal Usage Examples:
+    - To copy files to the trial directory on the storage filesystem:
+
+        pyarrow.fs.copy_files(
+            local_dir,
+            Path(storage.trial_fs_path, "subdir").as_posix(),
+            destination_filesystem=storage.filesystem
+        )
+
+    .. warning::
+        This is an experimental developer API and is subject to change
+        without notice between versions.
+    """
+
+    def __init__(
+        self,
+        storage_path: Union[str, os.PathLike],
+        experiment_dir_name: str,
+        sync_config: Optional[SyncConfig] = None,
+        storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
+        trial_dir_name: Optional[str] = None,
+        current_checkpoint_index: int = -1,
+    ):
+        from ray.tune.utils import date_str
+
+        self.custom_fs_provided = storage_filesystem is not None
+
+        # Invariant: (`storage_filesystem`, `storage_path`) is the location where
+        # *all* results can be accessed.
+        self.experiment_dir_name = experiment_dir_name
+        self.trial_dir_name = trial_dir_name
+        self.current_checkpoint_index = current_checkpoint_index
+        self.sync_config = sync_config or SyncConfig()
+
+        self.storage_filesystem, self.storage_fs_path = get_fs_and_path(
+            storage_path, storage_filesystem
+        )
+        self.storage_fs_path = Path(self.storage_fs_path).as_posix()
+
+        self.syncer: Syncer = _FilesystemSyncer(
+            storage_filesystem=self.storage_filesystem,
+            sync_period=self.sync_config.sync_period,
+            sync_timeout=self.sync_config.sync_timeout,
+        )
+
+        self._create_validation_file()
+        self._check_validation_file()
+
+        # Timestamp is used to create a unique session directory for the current
+        # training job. This is used to avoid conflicts when multiple training jobs
+        # run with the same name in the same cluster.
+        # This is set ONCE at the creation of the storage context, on the driver.
+        self._timestamp = date_str()
+
+    def __str__(self):
+        return (
+            "StorageContext<\n"
+            f"  storage_filesystem='{self.storage_filesystem.type_name}',\n"
+            f"  storage_fs_path='{self.storage_fs_path}',\n"
+            f"  experiment_dir_name='{self.experiment_dir_name}',\n"
+            f"  trial_dir_name='{self.trial_dir_name}',\n"
+            f"  current_checkpoint_index={self.current_checkpoint_index},\n"
+            ">"
+        )
+
+    def _create_validation_file(self):
+        """On the creation of a storage context, create a validation file at the
+        storage path to verify that the storage path can be written to.
+        This validation file is also used to check whether the storage path is
+        accessible by all nodes in the cluster."""
+        valid_file = Path(
+            self.experiment_fs_path, _VALIDATE_STORAGE_MARKER_FILENAME
+        ).as_posix()
+        self.storage_filesystem.create_dir(self.experiment_fs_path)
+        with self.storage_filesystem.open_output_stream(valid_file):
+            pass
+
+    def _check_validation_file(self):
+        """Checks that the validation file exists at the storage path."""
+        valid_file = Path(
+            self.experiment_fs_path, _VALIDATE_STORAGE_MARKER_FILENAME
+        ).as_posix()
+        if not _exists_at_fs_path(fs=self.storage_filesystem, fs_path=valid_file):
+            raise RuntimeError(
+                f"Unable to set up cluster storage with the following settings:\n{self}"
+                "\nCheck that all nodes in the cluster have read/write access "
+                "to the configured storage path. `RunConfig(storage_path)` should be "
+                "set to a cloud storage URI or a shared filesystem path accessible "
+                "by all nodes in your cluster ('s3://bucket' or '/mnt/nfs'). "
+                "A local path on the head node is not accessible by worker nodes. "
+                "See: https://docs.ray.io/en/latest/train/user-guides/persistent-storage.html"  # noqa: E501
+            )
+
+    def _update_checkpoint_index(self, metrics: Dict):
+        # Per default, increase by 1. This can be overwritten to customize checkpoint
+        # directories.
+        self.current_checkpoint_index += 1
+
+    def persist_current_checkpoint(self, checkpoint: "Checkpoint") -> "Checkpoint":
+        """Persists a given checkpoint to the current checkpoint path on the filesystem.
+
+        "Current" is defined by the `current_checkpoint_index` attribute of the
+        storage context.
+
+        This method copies the checkpoint files to the storage location.
+        It's up to the user to delete the original checkpoint files if desired.
+
+        For example, the original directory is typically a local temp directory.
+
+        Args:
+            checkpoint: The checkpoint to persist to (fs, checkpoint_fs_path).
+
+        Returns:
+            Checkpoint: A Checkpoint pointing to the persisted checkpoint location.
+        """
+        # TODO(justinvyu): Fix this cyclical import.
+
+        logger.debug(
+            "Copying checkpoint files to storage path:\n"
+            "({source_fs}, {source}) -> ({dest_fs}, {destination})".format(
+                source=checkpoint.path,
+                destination=self.checkpoint_fs_path,
+                source_fs=checkpoint.filesystem,
+                dest_fs=self.storage_filesystem,
+            )
+        )
+
+        # Raise an error if the storage path is not accessible when
+        # attempting to upload a checkpoint from a remote worker.
+        # Ex: If storage_path is a local path, then a validation marker
+        # will only exist on the head node but not the worker nodes.
+        self._check_validation_file()
+
+        self.storage_filesystem.create_dir(self.checkpoint_fs_path)
+        _pyarrow_fs_copy_files(
+            source=checkpoint.path,
+            destination=self.checkpoint_fs_path,
+            source_filesystem=checkpoint.filesystem,
+            destination_filesystem=self.storage_filesystem,
+        )
+
+        persisted_checkpoint = checkpoint.__class__(
+            filesystem=self.storage_filesystem,
+            path=self.checkpoint_fs_path,
+        )
+        logger.info(f"Checkpoint successfully created at: {persisted_checkpoint}")
+        return persisted_checkpoint
+
+    def persist_artifacts(self, force: bool = False) -> None:
+        """Persists all artifacts within `trial_local_dir` to storage.
+
+        This method possibly launches a background task to sync the trial dir,
+        depending on the `sync_period` + `sync_artifacts_on_checkpoint`
+        settings of `SyncConfig`.
+
+        `(local_fs, trial_working_dir) -> (storage_filesystem, trial_fs_path)`
+
+        Args:
+            force: If True, wait for a previous sync to finish, launch a new one,
+                and wait for that one to finish. By the end of a `force=True` call, the
+                latest version of the trial artifacts will be persisted.
+        """
+        if not self.sync_config.sync_artifacts:
+            return
+
+        # Skip if there are no artifacts to sync
+        is_empty = not any(os.scandir(self.trial_working_directory))
+        if is_empty:
+            return
+
+        if force:
+            self.syncer.wait()
+            self.syncer.sync_up(
+                local_dir=self.trial_working_directory, remote_dir=self.trial_fs_path
+            )
+            self.syncer.wait()
+        else:
+            self.syncer.sync_up_if_needed(
+                local_dir=self.trial_working_directory, remote_dir=self.trial_fs_path
+            )
+
+    @property
+    def experiment_fs_path(self) -> str:
+        """The path on the `storage_filesystem` to the experiment directory.
+
+        NOTE: This does not have a URI prefix anymore, since it has been stripped
+        by pyarrow.fs.FileSystem.from_uri already. The URI scheme information is
+        kept in `storage_filesystem` instead.
+        """
+        return Path(self.storage_fs_path, self.experiment_dir_name).as_posix()
+
+    def _get_session_path(self) -> str:
+        """The Ray Train/Tune session local directory used to stage files
+        before persisting to the storage filesystem."""
+        return Path(
+            _get_ray_train_session_dir(), self._timestamp, self.experiment_dir_name
+        ).as_posix()
+
+    @property
+    def experiment_driver_staging_path(self) -> str:
+        """The local filesystem path of the experiment directory on the driver node.
+
+        The driver is the node where `Trainer.fit`/`Tuner.fit` is being called.
+
+        This path is of the form:
+        `/tmp/ray/session_<session_id>/artifacts/<ray-train-job-timestamp>/
+        <experiment_dir_name>/driver_artifacts`
+
+        This should be used as the temporary staging location for files *on the driver*
+        before syncing them to `experiment_fs_path`.
+        For example, the search algorithm should dump its state to this directory.
+        See `trial_driver_staging_path` for writing trial-specific artifacts.
+
+        The directory is synced to
+        `{storage_path}/{experiment_dir_name}` periodically.
+        See `_ExperimentCheckpointManager.checkpoint` for where that happens.
+        """
+        return Path(self._get_session_path(), "driver_artifacts").as_posix()
+
+    @property
+    def trial_fs_path(self) -> str:
+        """The trial directory path on the `storage_filesystem`.
+
+        Raises a ValueError if `trial_dir_name` is not set beforehand.
+        """
+        if self.trial_dir_name is None:
+            raise RuntimeError(
+                "Should not access `trial_fs_path` without setting `trial_dir_name`"
+            )
+        return Path(self.experiment_fs_path, self.trial_dir_name).as_posix()
+
+    @property
+    def trial_driver_staging_path(self) -> str:
+        """The local filesystem path of the trial directory on the driver.
+
+        The driver is the node where `Trainer.fit`/`Tuner.fit` is being called.
+
+        This path is of the form:
+        `/tmp/ray/session_<session_id>/artifacts/<ray-train-job-timestamp>/
+        <experiment_dir_name>/driver_artifacts/<trial_dir_name>`
+
+        This should be used as the temporary location for files on the driver
+        before persisting them to `trial_fs_path`.
+
+        For example, callbacks (e.g., JsonLoggerCallback) should write trial-specific
+        logfiles within this directory.
+        """
+        if self.trial_dir_name is None:
+            raise RuntimeError(
+                "Should not access `trial_driver_staging_path` "
+                "without setting `trial_dir_name`"
+            )
+        return Path(self.experiment_driver_staging_path, self.trial_dir_name).as_posix()
+
+    @property
+    def trial_working_directory(self) -> str:
+        """The local filesystem path to trial working directory.
+
+        This path is of the form:
+        `/tmp/ray/session_<session_id>/artifacts/<ray-train-job-timestamp>/
+        <experiment_dir_name>/working_dirs/<trial_dir_name>`
+
+        Ray Train/Tune moves the remote actor's working directory to this path
+        by default, unless disabled by `RAY_CHDIR_TO_TRIAL_DIR` environment variable.
+
+        Writing files to this directory allows users to persist training artifacts
+        if `SyncConfig(sync_artifacts=True)` is set.
+        """
+        if self.trial_dir_name is None:
+            raise RuntimeError(
+                "Cannot access `trial_working_directory` without "
+                "setting `trial_dir_name`"
+            )
+        return Path(
+            self._get_session_path(), "working_dirs", self.trial_dir_name
+        ).as_posix()
+
+    @property
+    def checkpoint_fs_path(self) -> str:
+        """The current checkpoint directory path on the `storage_filesystem`.
+
+        "Current" refers to the checkpoint that is currently being created/persisted.
+        The user of this class is responsible for setting the `current_checkpoint_index`
+        (e.g., incrementing when needed).
+        """
+        return Path(self.trial_fs_path, self.checkpoint_dir_name).as_posix()
+
+    @property
+    def checkpoint_dir_name(self) -> str:
+        """The current checkpoint directory name, based on the checkpoint index."""
+        return StorageContext._make_checkpoint_dir_name(self.current_checkpoint_index)
+
+    @staticmethod
+    def get_experiment_dir_name(run_obj: Union[str, Callable, Type]) -> str:
+        from ray.tune.experiment import Experiment
+        from ray.tune.utils import date_str
+
+        run_identifier = Experiment.get_trainable_name(run_obj)
+
+        if bool(int(os.environ.get("TUNE_DISABLE_DATED_SUBDIR", 0))):
+            dir_name = run_identifier
+        else:
+            dir_name = "{}_{}".format(run_identifier, date_str())
+        return dir_name
+
+    @staticmethod
+    def _make_checkpoint_dir_name(index: int):
+        """Get the name of the checkpoint directory, given an index."""
+        return f"checkpoint_{index:06d}"

python/ray/train/_internal/syncer.py

@@ -0,0 +1,429 @@
+import abc
+import logging
+import threading
+import time
+import traceback
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from ray._private.thirdparty.tabulate.tabulate import tabulate
+from ray.util.annotations import Deprecated, DeveloperAPI
+from ray.widgets import Template
+
+logger = logging.getLogger(__name__)
+
+# Syncing period for syncing checkpoints between nodes or to cloud.
+DEFAULT_SYNC_PERIOD = 300
+
+# Default sync timeout after which syncing processes are aborted
+DEFAULT_SYNC_TIMEOUT = 1800
+
+
+@Deprecated
+@dataclass
+class SyncConfig:
+    sync_period: int = DEFAULT_SYNC_PERIOD
+    sync_timeout: int = DEFAULT_SYNC_TIMEOUT
+    sync_artifacts: bool = False
+    sync_artifacts_on_checkpoint: bool = True
+
+    def _repr_html_(self) -> str:
+        """Generate an HTML representation of the SyncConfig."""
+        return Template("scrollableTable.html.j2").render(
+            table=tabulate(
+                {
+                    "Setting": ["Sync period", "Sync timeout"],
+                    "Value": [self.sync_period, self.sync_timeout],
+                },
+                tablefmt="html",
+                showindex=False,
+                headers="keys",
+            ),
+            max_height="none",
+        )
+
+
+class _BackgroundProcess:
+    def __init__(self, fn: Callable):
+        self._fn = fn
+        self._process = None
+        self._result = {}
+        self._start_time = float("-inf")
+
+    @property
+    def is_running(self):
+        return self._process and self._process.is_alive()
+
+    @property
+    def start_time(self):
+        return self._start_time
+
+    def start(self, *args, **kwargs):
+        if self.is_running:
+            return False
+
+        self._result = {}
+
+        def entrypoint():
+            try:
+                result = self._fn(*args, **kwargs)
+            except Exception as e:
+                self._result["exception"] = e
+                return
+
+            self._result["result"] = result
+
+        self._process = threading.Thread(target=entrypoint)
+        self._process.daemon = True
+        self._process.start()
+        self._start_time = time.time()
+
+    def wait(self, timeout: Optional[float] = None) -> Any:
+        """Waits for the background process to finish running. Waits until the
+        background process has run for at least `timeout` seconds, counting from
+        the time when the process was started."""
+        if not self._process:
+            return None
+
+        time_remaining = None
+        if timeout:
+            elapsed = time.time() - self.start_time
+            time_remaining = max(timeout - elapsed, 0)
+
+        self._process.join(timeout=time_remaining)
+
+        if self._process.is_alive():
+            self._process = None
+            raise TimeoutError(
+                f"{getattr(self._fn, '__name__', str(self._fn))} did not finish "
+                f"running within the timeout of {timeout} seconds."
+            )
+
+        self._process = None
+
+        exception = self._result.get("exception")
+        if exception:
+            raise exception
+
+        result = self._result.get("result")
+
+        self._result = {}
+        return result
+
+
+@DeveloperAPI
+class Syncer(abc.ABC):
+    """Syncer class for synchronizing data between Ray nodes and remote (cloud) storage.
+
+    This class handles data transfer for two cases:
+
+    1. Synchronizing data such as experiment state snapshots from the driver to
+       cloud storage.
+    2. Synchronizing data such as trial checkpoints from remote trainables to
+       cloud storage.
+
+    Synchronizing tasks are usually asynchronous and can be awaited using ``wait()``.
+    The base class implements a ``wait_or_retry()`` API that will retry a failed
+    sync command.
+
+    The base class also exposes an API to only kick off syncs every ``sync_period``
+    seconds.
+
+    Args:
+        sync_period: The minimum time in seconds between sync operations, as
+            used by ``sync_up/down_if_needed``.
+        sync_timeout: The maximum time to wait for a sync process to finish before
+            issuing a new sync operation. Ex: should be used by ``wait`` if launching
+            asynchronous sync tasks.
+    """
+
+    def __init__(
+        self,
+        sync_period: float = DEFAULT_SYNC_PERIOD,
+        sync_timeout: float = DEFAULT_SYNC_TIMEOUT,
+    ):
+        self.sync_period = sync_period
+        self.sync_timeout = sync_timeout
+        self.last_sync_up_time = float("-inf")
+        self.last_sync_down_time = float("-inf")
+
+    @abc.abstractmethod
+    def sync_up(
+        self, local_dir: str, remote_dir: str, exclude: Optional[List] = None
+    ) -> bool:
+        """Synchronize local directory to remote directory.
+
+        This function can spawn an asynchronous process that can be awaited in
+        ``wait()``.
+
+        Args:
+            local_dir: Local directory to sync from.
+            remote_dir: Remote directory to sync up to. This is an URI
+                (``protocol://remote/path``).
+            exclude: Pattern of files to exclude, e.g.
+                ``["*/checkpoint_*]`` to exclude trial checkpoints.
+
+        Returns:
+            True if sync process has been spawned, False otherwise.
+
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def sync_down(
+        self, remote_dir: str, local_dir: str, exclude: Optional[List] = None
+    ) -> bool:
+        """Synchronize remote directory to local directory.
+
+        This function can spawn an asynchronous process that can be awaited in
+        ``wait()``.
+
+        Args:
+            remote_dir: Remote directory to sync down from. This is an URI
+                (``protocol://remote/path``).
+            local_dir: Local directory to sync to.
+            exclude: Pattern of files to exclude, e.g.
+                ``["*/checkpoint_*]`` to exclude trial checkpoints.
+
+        Returns:
+            True if sync process has been spawned, False otherwise.
+
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def delete(self, remote_dir: str) -> bool:
+        """Delete directory on remote storage.
+
+        This function can spawn an asynchronous process that can be awaited in
+        ``wait()``.
+
+        Args:
+            remote_dir: Remote directory to delete. This is an URI
+                (``protocol://remote/path``).
+
+        Returns:
+            True if sync process has been spawned, False otherwise.
+
+        """
+        raise NotImplementedError
+
+    def retry(self):
+        """Retry the last sync up, sync down, or delete command.
+
+        You should implement this method if you spawn asynchronous syncing
+        processes.
+        """
+        pass
+
+    def wait(self, timeout: Optional[float] = None):
+        """Wait for asynchronous sync command to finish.
+
+        You should implement this method if you spawn asynchronous syncing
+        processes. This method should timeout after the asynchronous command
+        has run for `sync_timeout` seconds and raise a `TimeoutError`.
+        """
+        pass
+
+    def sync_up_if_needed(
+        self, local_dir: str, remote_dir: str, exclude: Optional[List] = None
+    ) -> bool:
+        """Syncs up if time since last sync up is greater than sync_period.
+
+        Args:
+            local_dir: Local directory to sync from.
+            remote_dir: Remote directory to sync up to. This is an URI
+                (``protocol://remote/path``).
+            exclude: Pattern of files to exclude, e.g.
+                ``["*/checkpoint_*]`` to exclude trial checkpoints.
+
+        Returns:
+            The result of ``sync_up`` if a sync was triggered, otherwise ``None``.
+        """
+        now = time.time()
+        if now - self.last_sync_up_time >= self.sync_period:
+            result = self.sync_up(
+                local_dir=local_dir, remote_dir=remote_dir, exclude=exclude
+            )
+            self.last_sync_up_time = now
+            return result
+
+    def sync_down_if_needed(
+        self, remote_dir: str, local_dir: str, exclude: Optional[List] = None
+    ):
+        """Syncs down if time since last sync down is greater than sync_period.
+
+        Args:
+            remote_dir: Remote directory to sync down from. This is an URI
+                (``protocol://remote/path``).
+            local_dir: Local directory to sync to.
+            exclude: Pattern of files to exclude, e.g.
+                ``["*/checkpoint_*]`` to exclude trial checkpoints.
+
+        Returns:
+            The result of ``sync_down`` if a sync was triggered, otherwise ``None``.
+        """
+        now = time.time()
+        if now - self.last_sync_down_time >= self.sync_period:
+            result = self.sync_down(
+                remote_dir=remote_dir, local_dir=local_dir, exclude=exclude
+            )
+            self.last_sync_down_time = now
+            return result
+
+    def wait_or_retry(self, max_retries: int = 2, backoff_s: int = 5):
+        assert max_retries > 0
+        last_error_traceback = None
+        for i in range(max_retries + 1):
+            try:
+                self.wait()
+            except Exception as e:
+                attempts_remaining = max_retries - i
+
+                # If we're out of retries, then save the full traceback of the last
+                # error and show it when raising an exception.
+                if attempts_remaining == 0:
+                    last_error_traceback = traceback.format_exc()
+                    break
+
+                logger.error(
+                    f"The latest sync operation failed with the following error: "
+                    f"{repr(e)}\n"
+                    f"Retrying {attempts_remaining} more time(s) after sleeping "
+                    f"for {backoff_s} seconds..."
+                )
+                time.sleep(backoff_s)
+                self.retry()
+                continue
+            # Succeeded!
+            return
+        raise RuntimeError(
+            f"Failed sync even after {max_retries} retries. "
+            f"The latest sync failed with the following error:\n{last_error_traceback}"
+        )
+
+    def reset(self):
+        self.last_sync_up_time = float("-inf")
+        self.last_sync_down_time = float("-inf")
+
+    def close(self):
+        pass
+
+    def _repr_html_(self) -> str:
+        return
+
+
+class _BackgroundSyncer(Syncer):
+    """Syncer using a background process for asynchronous file transfer."""
+
+    def __init__(
+        self,
+        sync_period: float = DEFAULT_SYNC_PERIOD,
+        sync_timeout: float = DEFAULT_SYNC_TIMEOUT,
+    ):
+        super(_BackgroundSyncer, self).__init__(
+            sync_period=sync_period, sync_timeout=sync_timeout
+        )
+        self._sync_process = None
+        self._current_cmd = None
+
+    def _should_continue_existing_sync(self):
+        """Returns whether a previous sync is still running within the timeout."""
+        return (
+            self._sync_process
+            and self._sync_process.is_running
+            and time.time() - self._sync_process.start_time < self.sync_timeout
+        )
+
+    def _launch_sync_process(self, sync_command: Tuple[Callable, Dict]):
+        """Waits for the previous sync process to finish,
+        then launches a new process that runs the given command."""
+        if self._sync_process:
+            try:
+                self.wait()
+            except Exception:
+                logger.warning(
+                    f"Last sync command failed with the following error:\n"
+                    f"{traceback.format_exc()}"
+                )
+
+        self._current_cmd = sync_command
+        self.retry()
+
+    def sync_up(
+        self, local_dir: str, remote_dir: str, exclude: Optional[List] = None
+    ) -> bool:
+        if self._should_continue_existing_sync():
+            logger.debug(
+                f"Last sync still in progress, "
+                f"skipping sync up of {local_dir} to {remote_dir}"
+            )
+            return False
+
+        sync_up_cmd = self._sync_up_command(
+            local_path=local_dir, uri=remote_dir, exclude=exclude
+        )
+        self._launch_sync_process(sync_up_cmd)
+
+        return True
+
+    def _sync_up_command(
+        self, local_path: str, uri: str, exclude: Optional[List] = None
+    ) -> Tuple[Callable, Dict]:
+        raise NotImplementedError
+
+    def sync_down(
+        self, remote_dir: str, local_dir: str, exclude: Optional[List] = None
+    ) -> bool:
+        if self._should_continue_existing_sync():
+            logger.warning(
+                f"Last sync still in progress, "
+                f"skipping sync down of {remote_dir} to {local_dir}"
+            )
+            return False
+
+        sync_down_cmd = self._sync_down_command(uri=remote_dir, local_path=local_dir)
+        self._launch_sync_process(sync_down_cmd)
+
+        return True
+
+    def _sync_down_command(self, uri: str, local_path: str) -> Tuple[Callable, Dict]:
+        raise NotImplementedError
+
+    def delete(self, remote_dir: str) -> bool:
+        if self._should_continue_existing_sync():
+            logger.warning(
+                f"Last sync still in progress, skipping deletion of {remote_dir}"
+            )
+            return False
+
+        delete_cmd = self._delete_command(uri=remote_dir)
+        self._launch_sync_process(delete_cmd)
+
+        return True
+
+    def _delete_command(self, uri: str) -> Tuple[Callable, Dict]:
+        raise NotImplementedError
+
+    def wait(self, timeout: Optional[float] = None):
+        if self._sync_process:
+            try:
+                self._sync_process.wait(timeout=timeout or self.sync_timeout)
+            except Exception as e:
+                raise e
+            finally:
+                # Regardless of whether the sync process succeeded within the timeout,
+                # clear the sync process so a new one can be created.
+                self._sync_process = None
+
+    def retry(self):
+        if not self._current_cmd:
+            raise RuntimeError("No sync command set, cannot retry.")
+        cmd, kwargs = self._current_cmd
+        self._sync_process = _BackgroundProcess(cmd)
+        self._sync_process.start(**kwargs)
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["_sync_process"] = None
+        return state

python/ray/train/_internal/utils.py

@@ -0,0 +1,230 @@
+import abc
+import functools
+import inspect
+import logging
+import os
+import socket
+from typing import (
+    Any,
+    Callable,
+    ContextManager,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+import ray
+from ray._common.network_utils import find_free_port, is_ipv6
+from ray.actor import ActorHandle
+from ray.air._internal.util import (
+    StartTraceback,
+    StartTracebackWithWorkerRank,
+)
+from ray.exceptions import RayActorError
+from ray.types import ObjectRef
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+def check_for_failure(
+    remote_values: List[ObjectRef],
+) -> Tuple[bool, Optional[Exception]]:
+    """Check for actor failure when retrieving the remote values.
+
+    Args:
+        remote_values: List of object references from Ray actor methods.
+
+    Returns:
+        A tuple of (bool, Exception). The bool is
+        True if evaluating all object references is successful, False otherwise.
+    """
+    unfinished = remote_values.copy()
+
+    while len(unfinished) > 0:
+        finished, unfinished = ray.wait(unfinished)
+
+        # If a failure occurs the ObjectRef will be marked as finished.
+        # Calling ray.get will expose the failure as a RayActorError.
+        for object_ref in finished:
+            # Everything in finished has either failed or completed
+            # successfully.
+            try:
+                ray.get(object_ref)
+            except RayActorError as exc:
+                failed_actor_rank = remote_values.index(object_ref)
+                logger.info(f"Worker {failed_actor_rank} has failed.")
+                return False, exc
+            except Exception as exc:
+                # Other (e.g. training) errors should be directly raised
+                failed_worker_rank = remote_values.index(object_ref)
+                raise StartTracebackWithWorkerRank(
+                    worker_rank=failed_worker_rank
+                ) from exc
+
+    return True, None
+
+
+def get_address_and_port() -> Tuple[str, int]:
+    """Returns the IP address and a free port on this node."""
+    addr = ray.util.get_node_ip_address()
+    port = find_free_port(socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET)
+    return addr, port
+
+
+def update_env_vars(env_vars: Dict[str, Any]):
+    """Updates the environment variables on this worker process.
+
+    Args:
+        env_vars: Environment variables to set.
+    """
+    sanitized = {k: str(v) for k, v in env_vars.items()}
+    os.environ.update(sanitized)
+
+
+def count_required_parameters(fn: Callable) -> int:
+    """Counts the number of required parameters of a function.
+
+    NOTE: *args counts as 1 required parameter.
+
+    Args:
+        fn: The function whose required parameters should be counted.
+
+    Returns:
+        The number of required parameters of ``fn``.
+
+    Examples:
+
+        >>> def fn(a, b, /, c, *args, d=1, e=2, **kwargs):
+        ...    pass
+        >>> count_required_parameters(fn)
+        4
+
+        >>> fn = lambda: 1
+        >>> count_required_parameters(fn)
+        0
+
+        >>> def fn(config, a, b=1, c=2):
+        ...     pass
+        >>> from functools import partial
+        >>> count_required_parameters(partial(fn, a=0))
+        1
+    """
+    params = inspect.signature(fn).parameters.values()
+
+    positional_param_kinds = {
+        inspect.Parameter.POSITIONAL_ONLY,
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        inspect.Parameter.VAR_POSITIONAL,
+    }
+    return len(
+        [
+            p
+            for p in params
+            if p.default == inspect.Parameter.empty and p.kind in positional_param_kinds
+        ]
+    )
+
+
+def construct_train_func(
+    train_func: Union[Callable[[], T], Callable[[Dict[str, Any]], T]],
+    config: Optional[Dict[str, Any]],
+    train_func_context: ContextManager,
+    fn_arg_name: Optional[str] = "train_func",
+    discard_returns: bool = False,
+) -> Callable[[], T]:
+    """Validates and constructs the training function to execute.
+
+    Args:
+        train_func: The training function to execute.
+            This can either take in no arguments or a ``config`` dict.
+        config: Configurations to pass into ``train_func``. If None then an empty
+            Dict will be created.
+        train_func_context: Context manager for user's `train_func`, which executes
+            backend-specific logic before and after the training function.
+        fn_arg_name: The name of training function to use for error messages.
+        discard_returns: Whether to discard any returns from train_func or not.
+
+    Returns:
+        A valid training function.
+
+    Raises:
+        ValueError: if the input ``train_func`` is invalid.
+    """
+    num_required_params = count_required_parameters(train_func)
+
+    if discard_returns:
+        # Discard any returns from the function so that
+        # BackendExecutor doesn't try to deserialize them.
+        # Those returns are inaccesible with AIR anyway.
+        @functools.wraps(train_func)
+        def discard_return_wrapper(*args, **kwargs):
+            try:
+                train_func(*args, **kwargs)
+            except Exception as e:
+                raise StartTraceback from e
+
+        wrapped_train_func = discard_return_wrapper
+    else:
+        wrapped_train_func = train_func
+
+    if num_required_params > 1:
+        err_msg = (
+            f"{fn_arg_name} should take in 0 or 1 required arguments, but it accepts "
+            f"{num_required_params} required arguments instead."
+        )
+        raise ValueError(err_msg)
+    elif num_required_params == 1:
+        config = {} if config is None else config
+
+        @functools.wraps(wrapped_train_func)
+        def train_fn():
+            try:
+                with train_func_context():
+                    return wrapped_train_func(config)
+            except Exception as e:
+                raise StartTraceback from e
+
+    else:  # num_params == 0
+
+        @functools.wraps(wrapped_train_func)
+        def train_fn():
+            try:
+                with train_func_context():
+                    return wrapped_train_func()
+            except Exception as e:
+                raise StartTraceback from e
+
+    return train_fn
+
+
+class Singleton(abc.ABCMeta):
+    """Singleton Abstract Base Class
+
+    https://stackoverflow.com/questions/33364070/implementing
+    -singleton-as-metaclass-but-for-abstract-classes
+    """
+
+    _instances = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
+        return cls._instances[cls]
+
+
+class ActorWrapper:
+    """Wraps an actor to provide same API as using the base class directly."""
+
+    def __init__(self, actor: ActorHandle):
+        self.actor = actor
+
+    def __getattr__(self, item):
+        # The below will fail if trying to access an attribute (not a method) from the
+        # actor.
+        actor_method = getattr(self.actor, item)
+        return lambda *args, **kwargs: ray.get(actor_method.remote(*args, **kwargs))

python/ray/train/_internal/worker_group.py

@@ -0,0 +1,445 @@
+import copy
+import logging
+import os
+import socket
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
+
+import ray
+from ray.actor import ActorHandle
+from ray.air._internal.util import exception_cause, skip_exceptions
+from ray.train._internal.base_worker_group import BaseWorkerGroup
+from ray.types import ObjectRef
+from ray.util.placement_group import PlacementGroup
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+class RayTrainWorker:
+    """A class to execute arbitrary functions. Does not hold any state."""
+
+    def __execute(self, func: Callable[..., T], *args, **kwargs) -> T:
+        """Executes the input function and returns the output.
+
+        Args:
+            func: The function to execute.
+            *args: Positional arguments to pass into ``func``.
+            **kwargs: Keyword arguments to pass into ``func``.
+
+        Returns:
+            The result of calling ``func`` with the provided arguments.
+        """
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            skipped = skip_exceptions(e)
+            raise skipped from exception_cause(skipped)
+
+
+@dataclass
+class WorkerMetadata:
+    """Metadata for each worker/actor.
+
+    This information is expected to stay the same throughout the lifetime of
+    actor.
+
+    Args:
+        node_id: ID of the node this worker is on.
+        node_ip: IP address of the node this worker is on.
+        hostname: Hostname that this worker is on.
+        resource_ids: Map of accelerator resources
+            ("GPU", "neuron_cores", ..) to their IDs.
+        pid: Process ID of this worker.
+    """
+
+    node_id: str
+    node_ip: str
+    hostname: str
+    resource_ids: Dict[str, List[str]]
+    pid: int
+
+
+@dataclass
+class Worker:
+    """Class representing a Worker."""
+
+    actor: ActorHandle
+    metadata: WorkerMetadata
+
+
+def create_executable_class(executable_cls: Optional[Type] = None) -> Type:
+    """Create the executable class to use as the Ray actors."""
+    if not executable_cls:
+        return RayTrainWorker
+    elif issubclass(executable_cls, RayTrainWorker):
+        return executable_cls
+    else:
+
+        class _WrappedExecutable(executable_cls, RayTrainWorker):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+
+        return _WrappedExecutable
+
+
+def construct_metadata() -> WorkerMetadata:
+    """Creates metadata for this worker.
+
+    This function is expected to be run on the actor.
+    """
+    node_id = ray.get_runtime_context().get_node_id()
+    node_ip = ray.util.get_node_ip_address()
+    hostname = socket.gethostname()
+    accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
+    pid = os.getpid()
+
+    return WorkerMetadata(
+        node_id=node_id,
+        node_ip=node_ip,
+        hostname=hostname,
+        resource_ids=accelerator_ids,
+        pid=pid,
+    )
+
+
+class WorkerGroup(BaseWorkerGroup):
+    """Group of Ray Actors that can execute arbitrary functions.
+
+    ``WorkerGroup`` launches Ray actors according to the given
+    specification. It can then execute arbitrary Python functions in each of
+    these workers.
+
+    If not enough resources are available to launch the actors, the Ray
+    cluster will automatically scale up if autoscaling is enabled.
+
+    Args:
+        num_workers: The number of workers (Ray actors) to launch.
+            Defaults to 1.
+        resources_per_worker: Dictionary specifying the resources that will be
+            requested for each worker. Defaults to {"CPU": 1}.
+        actor_cls: If specified use this class as the remote actors.
+        actor_cls_args: If ``actor_cls`` is provided, these positional args will
+            be used for the worker initialization.
+        actor_cls_kwargs: If ``actor_cls`` is provided, these keyword args will
+            be used for the worker initialization.
+        placement_group: The placement group that workers
+            should be created in. Defaults to "default" which will inherit the
+            parent placement group (if child tasks should be captured).
+
+
+    Example:
+
+    .. code_block:: python
+
+        worker_group = WorkerGroup(num_workers=2)
+        output = worker_group.execute(lambda: 1)
+        assert len(output) == 2
+        assert all(o == 1 for o in output)
+    """
+
+    def __init__(
+        self,
+        num_workers: int = 1,
+        resources_per_worker: Optional[Dict[str, float]] = None,
+        actor_cls: Type = None,
+        actor_cls_args: Optional[Tuple] = None,
+        actor_cls_kwargs: Optional[Dict] = None,
+        placement_group: Union[PlacementGroup, str] = "default",
+    ):
+        if resources_per_worker is None:
+            resources_per_worker = {"CPU": 1}
+        else:
+            resources_per_worker = resources_per_worker.copy()
+
+        if num_workers <= 0:
+            raise ValueError(
+                "The provided `num_workers` must be greater "
+                f"than 0. Received num_workers={num_workers} "
+                f"instead."
+            )
+
+        if any(v < 0 for v in resources_per_worker.values()):
+            raise ValueError(
+                "The number of resources per worker must not be negative. "
+                f"Received resources_per_worker={resources_per_worker}."
+            )
+
+        if (actor_cls_args or actor_cls_kwargs) and not actor_cls:
+            raise ValueError(
+                "`actor_cls_args` or `actor_class_kwargs` are "
+                "passed in but no `actor_cls` is passed in."
+            )
+
+        self.num_workers = num_workers
+        self.resources_per_worker = resources_per_worker
+
+        _resources_per_worker = copy.deepcopy(resources_per_worker)
+        self.num_cpus_per_worker = _resources_per_worker.pop("CPU", 0)
+        self.num_gpus_per_worker = _resources_per_worker.pop("GPU", 0)
+        self.memory_per_worker = _resources_per_worker.pop("memory", 0)
+        self.workers = []
+        self._base_cls = create_executable_class(actor_cls)
+        assert issubclass(self._base_cls, RayTrainWorker)
+
+        self._actor_cls_args = actor_cls_args or []
+        self._actor_cls_kwargs = actor_cls_kwargs or {}
+
+        self._placement_group = placement_group
+
+        # TODO(matt): Validate resources. Fast-fail if it is impossible to
+        #  handle the request, rather than hang indefinitely.
+        self._remote_cls = ray.remote(
+            num_cpus=self.num_cpus_per_worker,
+            num_gpus=self.num_gpus_per_worker,
+            memory=self.memory_per_worker,
+            resources=_resources_per_worker,
+        )(self._base_cls)
+        self.start()
+
+    def start(self):
+        """Starts all the workers in this worker group."""
+        if self.workers and len(self.workers) > 0:
+            raise RuntimeError(
+                "The workers have already been started. "
+                "Please call `shutdown` first if you want to "
+                "restart them."
+            )
+
+        logger.debug(f"Starting {self.num_workers} workers.")
+        self.add_workers(self.num_workers)
+        logger.debug(f"{len(self.workers)} workers have successfully started.")
+
+    def shutdown(self, patience_s: float = 5):
+        """Shutdown all the workers in this worker group.
+
+        Args:
+            patience_s: Attempt a graceful shutdown
+                of the workers for this many seconds. Fallback to force kill
+                if graceful shutdown is not complete after this time. If
+                this is less than or equal to 0, immediately force kill all
+                workers.
+        """
+        logger.debug(f"Shutting down {len(self.workers)} workers.")
+        if patience_s <= 0:
+            for worker in self.workers:
+                ray.kill(worker.actor)
+        else:
+            done_refs = [w.actor.__ray_terminate__.remote() for w in self.workers]
+            # Wait for actors to die gracefully.
+            done, not_done = ray.wait(done_refs, timeout=patience_s)
+            if not_done:
+                logger.debug("Graceful termination failed. Falling back to force kill.")
+                # If all actors are not able to die gracefully, then kill them.
+                for worker in self.workers:
+                    ray.kill(worker.actor)
+
+        logger.debug("Shutdown successful.")
+        self.workers = []
+
+    def execute_async(self, func: Callable[..., T], *args, **kwargs) -> List[ObjectRef]:
+        """Execute ``func`` on each worker and return the futures.
+
+        Args:
+            func: A function to call on each worker.
+            *args: Positional arguments passed directly into ``func``.
+            **kwargs: Keyword arguments passed directly into ``func``.
+
+        Returns:
+            (List[ObjectRef]) A list of ``ObjectRef`` representing the
+                output of ``func`` from each worker. The order is the same
+                as ``self.workers``.
+
+        """
+        if len(self.workers) <= 0:
+            raise RuntimeError(
+                "There are no active workers. This worker "
+                "group has most likely been shut down. Please"
+                "create a new WorkerGroup or restart this one."
+            )
+
+        return [
+            w.actor._RayTrainWorker__execute.options(
+                name=f"_RayTrainWorker__execute.{func.__name__}"
+            ).remote(func, *args, **kwargs)
+            for w in self.workers
+        ]
+
+    def execute(self, func: Callable[..., T], *args, **kwargs) -> List[T]:
+        """Execute ``func`` on each worker and return the outputs of ``func``.
+
+        Args:
+            func: A function to call on each worker.
+            *args: Positional arguments passed directly into ``func``.
+            **kwargs: Keyword arguments passed directly into ``func``.
+
+        Returns:
+            (List[T]) A list containing the output of ``func`` from each
+                worker. The order is the same as ``self.workers``.
+
+        """
+        # TODO: Add a timeout in the case of a hang, particularly
+        # relevant when func is TorchConfig.on_shutdown
+        return ray.get(self.execute_async(func, *args, **kwargs))
+
+    def execute_single_async(
+        self, worker_index: int, func: Callable[..., T], *args, **kwargs
+    ) -> ObjectRef:
+        """Execute ``func`` on worker ``worker_index`` and return futures.
+
+        Args:
+            worker_index: The index to execute func on.
+            func: A function to call on the first worker.
+            *args: Positional arguments passed directly into ``func``.
+            **kwargs: Keyword arguments passed directly into ``func``.
+
+        Returns:
+            (ObjectRef) An ObjectRef representing the output of func.
+
+        """
+        if worker_index >= len(self.workers):
+            raise ValueError(
+                f"The provided worker_index {worker_index} is "
+                f"not valid for {self.num_workers} workers."
+            )
+        return (
+            self.workers[worker_index]
+            .actor._RayTrainWorker__execute.options(
+                name=f"_RayTrainWorker__execute.{func.__name__}"
+            )
+            .remote(func, *args, **kwargs)
+        )
+
+    def execute_single(
+        self, worker_index: int, func: Callable[..., T], *args, **kwargs
+    ) -> T:
+        """Execute ``func`` on worker with index ``worker_index``.
+
+        Args:
+            worker_index: The index to execute func on.
+            func: A function to call on the first worker.
+            *args: Positional arguments passed directly into ``func``.
+            **kwargs: Keyword arguments passed directly into ``func``.
+
+        Returns:
+            (T) The output of func.
+
+        """
+
+        return ray.get(self.execute_single_async(worker_index, func, *args, **kwargs))
+
+    def remove_workers(self, worker_indexes: List[int]):
+        """Removes the workers with the specified indexes.
+
+        The removed workers will go out of scope and their actor processes
+        will be terminated.
+
+        Args:
+            worker_indexes: The indexes of the workers to remove.
+        """
+        new_workers = []
+        for i in range(len(self.workers)):
+            if i not in worker_indexes:
+                new_workers.append(self.workers[i])
+        self.workers = new_workers
+
+    def add_workers(self, num_workers: int):
+        """Adds ``num_workers`` to this WorkerGroup.
+
+        Note: Adding workers when the cluster/placement group is at capacity
+        may lead to undefined hanging behavior. If you are attempting to
+        replace existing workers in the WorkerGroup, remove_workers() should
+        be called first.
+
+        Args:
+            num_workers: The number of workers to add.
+        """
+        new_actors = []
+        new_actor_metadata = []
+        for _ in range(num_workers):
+            actor = self._remote_cls.options(
+                placement_group=self._placement_group
+            ).remote(*self._actor_cls_args, **self._actor_cls_kwargs)
+            new_actors.append(actor)
+            new_actor_metadata.append(
+                actor._RayTrainWorker__execute.options(
+                    name="_RayTrainWorker__execute.construct_metadata"
+                ).remote(construct_metadata)
+            )
+
+        # Get metadata from all actors.
+        metadata = ray.get(new_actor_metadata)
+
+        for i in range(len(new_actors)):
+            self.workers.append(Worker(actor=new_actors[i], metadata=metadata[i]))
+
+    def sort_workers_by_node_id_and_gpu_id(self, _first_node_id: Optional[str] = None):
+        """Reorder the workers by their node id and the lowest GPU id.
+
+        This is useful for collocating workers on the same node.
+
+        Example:
+            Given workers with the following attributes:
+                worker_0: node_id=1, gpu_ids=[1]
+                worker_1: node_id=0, gpu_ids=[0]
+                worker_2: node_id=1, gpu_ids=[0]
+                worker_3: node_id=0, gpu_ids=[1]
+
+            The function will perform the following steps:
+                1. Group by node ID:
+                    node_id=0: worker_1, worker_3
+                    node_id=1: worker_0, worker_2
+
+                2. Sort each group by GPU ID:
+                    node_id=0: worker_1 (gpu_id=0), worker_3 (gpu_id=1)
+                    node_id=1: worker_2 (gpu_id=0), worker_0 (gpu_id=1)
+
+            Resulting in the order: [worker_1, worker_3, worker_2, worker_0]
+
+        Args:
+            _first_node_id: The first ID to group by.
+                Set this to the node ID of the trainer coordinator to ensure that the
+                rank 0 worker is on the same node, allowing additional resources to
+                be specified for rank 0 workers via
+                `ScalingConfig(trainer_resources=)`.
+        """
+        node_id_to_workers = defaultdict(list)
+
+        if _first_node_id is not None:
+            node_id_to_workers[_first_node_id] = []
+
+        for worker in self.workers:
+            node_id_to_workers[worker.metadata.node_id].append(worker)
+
+        # Sort workers on the same node by the lowest GPU id
+        # More details: https://github.com/ray-project/ray/issues/40803
+        def get_lowest_gpu_id(worker) -> int:
+            gpu_ids = worker.metadata.resource_ids.get("GPU", [])
+            # If there are no GPU IDs, return 0 as a default
+            if not gpu_ids:
+                return 0
+
+            # Attempt to convert GPU IDs to integers and find the minimum ID.
+            # Fallback to return the minimum string-based ID
+            try:
+                return min(int(gpu_id) for gpu_id in gpu_ids)
+            except ValueError:
+                return min(gpu_ids)
+
+        for node_id in node_id_to_workers:
+            node_id_to_workers[node_id].sort(key=get_lowest_gpu_id)
+
+        sorted_workers = []
+        for workers in node_id_to_workers.values():
+            sorted_workers.extend(workers)
+
+        self.workers = sorted_workers
+
+    def __len__(self):
+        return len(self.workers)
+
+    def get_resources_per_worker(self) -> dict:
+        """Get the resources allocated per worker."""
+        return copy.deepcopy(self.resources_per_worker)

python/ray/train/backend.py

@@ -0,0 +1,74 @@
+import logging
+from contextlib import nullcontext
+from typing import Any, Dict, TypeVar
+
+from ray.train._internal.base_worker_group import BaseWorkerGroup
+from ray.train._internal.utils import Singleton
+from ray.util.annotations import DeveloperAPI
+from ray.widgets import make_table_html_repr
+
+EncodedData = TypeVar("EncodedData")
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class BackendConfig:
+    """Parent class for configurations of training backend."""
+
+    @property
+    def backend_cls(self):
+        return Backend
+
+    @property
+    def train_func_context(self):
+        return nullcontext
+
+    @property
+    def framework(self):
+        return None
+
+    def _repr_html_(self) -> str:
+        return make_table_html_repr(obj=self, title=type(self).__name__)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Returns serializable dictionary representation of the backend config.
+        Subclasses can override this to expose framework-specific configuration.
+
+        The fields here are used for state export of the backend config.
+        If a field is not serializable, it should be excluded.
+        """
+        return {}
+
+
+@DeveloperAPI
+class Backend(metaclass=Singleton):
+    """Singleton for distributed communication backend.
+
+    Attributes:
+        share_cuda_visible_devices: If True, each worker
+            process will have CUDA_VISIBLE_DEVICES set as the visible device
+            IDs of all workers on the same node for this training instance.
+            If False, each worker will have CUDA_VISIBLE_DEVICES set to the
+            device IDs allocated by Ray for that worker.
+    """
+
+    share_cuda_visible_devices: bool = False
+    has_replica_groups: bool = False
+
+    def on_start(self, worker_group: BaseWorkerGroup, backend_config: BackendConfig):
+        """Logic for starting this backend."""
+        pass
+
+    def on_shutdown(self, worker_group: BaseWorkerGroup, backend_config: BackendConfig):
+        """Logic for shutting down the backend."""
+        pass
+
+    def on_training_start(
+        self, worker_group: BaseWorkerGroup, backend_config: BackendConfig
+    ):
+        """Logic ran right before training is started.
+
+        Session API is available at this point."""
+        pass

python/ray/train/base_trainer.py

@@ -0,0 +1,954 @@
+import abc
+import copy
+import inspect
+import json
+import logging
+import os
+import warnings
+from functools import partial
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
+
+import pyarrow.fs
+
+import ray
+import ray.cloudpickle as pickle
+from ray._common.usage import usage_lib
+from ray._private.dict import deep_update
+from ray.air._internal import usage as air_usage
+from ray.air._internal.config import ensure_only_allowed_dataclass_keys_updated
+from ray.air._internal.usage import AirEntrypoint
+from ray.air.config import RunConfig, ScalingConfig
+from ray.air.result import Result
+from ray.train import Checkpoint
+from ray.train._internal.session import get_session
+from ray.train._internal.storage import (
+    StorageContext,
+    _exists_at_fs_path,
+    get_fs_and_path,
+)
+from ray.train.constants import (
+    V2_MIGRATION_GUIDE_MESSAGE,
+    _v2_migration_warnings_enabled,
+)
+from ray.train.context import _GET_METADATA_DEPRECATION_MESSAGE
+from ray.train.utils import _log_deprecation_warning
+from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data import Dataset
+    from ray.tune import Trainable
+
+_TRAINER_PKL = "trainer.pkl"
+
+# A type representing either a ray.data.Dataset or a function that returns a
+# ray.data.Dataset and accepts no arguments.
+GenDataset = Union["Dataset", Callable[[], "Dataset"]]
+
+
+logger = logging.getLogger(__name__)
+
+PREPROCESSOR_DEPRECATION_MESSAGE = (
+    "The `preprocessor` argument to Trainers is deprecated as of Ray 2.7. "
+    "Instead, use the Preprocessor `fit` and `transform` APIs directly on the Ray "
+    "Dataset. For any state that needs to be saved to the trained checkpoint, pass it "
+    "in using the `metadata` argument of the `Trainer`. "
+    "For a full example, see "
+    "https://docs.ray.io/en/master/train/user-guides/data-loading-preprocessing.html#preprocessing-structured-data "  # noqa:E501
+)
+
+_TRAINER_RESTORE_DEPRECATION_WARNING = (
+    "The `restore` and `can_restore` APIs are deprecated and "
+    f"will be removed in a future release. {V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING = (
+    "`resume_from_checkpoint` is deprecated and will be removed in an upcoming "
+    f"release. {V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+
+@PublicAPI(stability="beta")
+class TrainingFailedError(RuntimeError):
+    """An error indicating that training has failed."""
+
+    _RESTORE_MSG = (
+        "The Ray Train run failed. Please inspect the previous error messages for a "
+        "cause. After fixing the issue (assuming that the error is not caused by "
+        "your own application logic, but rather an error such as OOM), you can restart "
+        "the run from scratch or continue this run.\n"
+        "To continue this run, you can use: "
+        '`trainer = {trainer_cls_name}.restore("{path}")`.'
+    )
+
+    _FAILURE_CONFIG_MSG = (
+        "To start a new run that will retry on training failures, set "
+        "`train.RunConfig(failure_config=train.FailureConfig(max_failures))` "
+        "in the Trainer's `run_config` with `max_failures > 0`, or `max_failures = -1` "
+        "for unlimited retries."
+    )
+
+
+def _train_coordinator_fn(
+    config: dict, trainer_cls: Type["BaseTrainer"], metadata: dict
+):
+    """This is the function that defines the logic of the Ray Train coordinator.
+    This is responsible for setting up a remote instance of the `trainer_cls`
+    (a different instance than the one calling `trainer.fit` on the driver!)
+    and running the training loop.
+    """
+    assert metadata is not None, metadata
+    # Propagate user metadata from the Trainer constructor.
+    get_session().metadata = metadata
+
+    # config already contains merged values.
+    # Instantiate new Trainer in Trainable.
+    trainer = trainer_cls(**config)
+
+    # Get the checkpoint from Tune and pass it to workers later on.
+    checkpoint = ray.tune.get_checkpoint()
+    if checkpoint:
+        # Set `starting_checkpoint` for auto-recovery fault-tolerance
+        # as well as manual restoration.
+        trainer.starting_checkpoint = checkpoint
+    # else: Train will restore from the user-provided
+    # `resume_from_checkpoint` == `starting_checkpoint`.
+
+    # Evaluate datasets if they are wrapped in a factory.
+    trainer.datasets = {
+        k: d() if callable(d) else d for k, d in trainer.datasets.items()
+    }
+
+    trainer.setup()
+    trainer.training_loop()
+
+
+@DeveloperAPI
+class BaseTrainer(abc.ABC):
+    """Defines interface for distributed training on Ray.
+
+    Note: The base ``BaseTrainer`` class cannot be instantiated directly. Only
+    one of its subclasses can be used.
+
+    Note to developers: If a new trainer is added, please update
+    `air/_internal/usage.py`.
+
+    **How does a trainer work?**
+
+    - First, initialize the Trainer. The initialization runs locally,
+      so heavyweight setup should not be done in ``__init__``.
+    - Then, when you call ``trainer.fit()``, the Trainer is serialized
+      and copied to a remote Ray actor. The following methods are then
+      called in sequence on the remote actor.
+    - ``trainer.setup()``: Any heavyweight Trainer setup should be
+      specified here.
+    - ``trainer.training_loop()``: Executes the main training logic.
+    - Calling ``trainer.fit()`` will return a ``ray.result.Result``
+      object where you can access metrics from your training run, as well
+      as any checkpoints that may have been saved.
+
+    **How do I create a new Trainer?**
+
+    Subclass ``ray.train.trainer.BaseTrainer``, and override the ``training_loop``
+    method, and optionally ``setup``.
+
+    .. testcode::
+        :skipif: True
+
+        import torch
+
+        from ray.train.trainer import BaseTrainer
+        from ray import train, tune
+
+
+        class MyPytorchTrainer(BaseTrainer):
+            def setup(self):
+                self.model = torch.nn.Linear(1, 1)
+                self.optimizer = torch.optim.SGD(
+                    self.model.parameters(), lr=0.1)
+
+            def training_loop(self):
+                # You can access any Trainer attributes directly in this method.
+                # self.datasets["train"] has already been
+                dataset = self.datasets["train"]
+
+                torch_ds = dataset.iter_torch_batches(dtypes=torch.float)
+                loss_fn = torch.nn.MSELoss()
+
+                for epoch_idx in range(10):
+                    loss = 0
+                    num_batches = 0
+                    torch_ds = dataset.iter_torch_batches(
+                        dtypes=torch.float, batch_size=2
+                    )
+                    for batch in torch_ds:
+                        X = torch.unsqueeze(batch["x"], 1)
+                        y = torch.unsqueeze(batch["y"], 1)
+                        # Compute prediction error
+                        pred = self.model(X)
+                        batch_loss = loss_fn(pred, y)
+
+                        # Backpropagation
+                        self.optimizer.zero_grad()
+                        batch_loss.backward()
+                        self.optimizer.step()
+
+                        loss += batch_loss.item()
+                        num_batches += 1
+                    loss /= num_batches
+
+                    # Use Tune functions to report intermediate
+                    # results.
+                    train.report({"loss": loss, "epoch": epoch_idx})
+
+
+        # Initialize the Trainer, and call Trainer.fit()
+        import ray
+        train_dataset = ray.data.from_items(
+            [{"x": i, "y": i} for i in range(10)])
+        my_trainer = MyPytorchTrainer(datasets={"train": train_dataset})
+        result = my_trainer.fit()
+
+    Args:
+        scaling_config: Configuration for how to scale training.
+        run_config: Configuration for the execution of the training run.
+        datasets: Any Datasets to use for training. Use the key "train"
+            to denote which dataset is the training dataset.
+        metadata: Dict that should be made available via
+            `train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
+            for checkpoints saved from this Trainer. Must be JSON-serializable.
+        resume_from_checkpoint: A checkpoint to resume training from.
+    """
+
+    _scaling_config_allowed_keys: List[str] = [
+        "trainer_resources",
+    ]
+    _handles_checkpoint_freq: bool = False
+    _handles_checkpoint_at_end: bool = False
+
+    # fields to propagate to Tuner param_space.
+    # See `BaseTrainer._extract_fields_for_tuner_param_space` for more details.
+    _fields_for_tuner_param_space = []
+
+    def __init__(
+        self,
+        *,
+        scaling_config: Optional[ScalingConfig] = None,
+        run_config: Optional[RunConfig] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        resume_from_checkpoint: Optional[Checkpoint] = None,
+    ):
+        self.scaling_config = (
+            scaling_config if scaling_config is not None else ScalingConfig()
+        )
+        self.run_config = (
+            copy.copy(run_config) if run_config is not None else RunConfig()
+        )
+        self.metadata = metadata
+        self.datasets = datasets if datasets is not None else {}
+        self.starting_checkpoint = resume_from_checkpoint
+
+        if _v2_migration_warnings_enabled():
+            if metadata is not None:
+                _log_deprecation_warning(_GET_METADATA_DEPRECATION_MESSAGE)
+            if resume_from_checkpoint is not None:
+                _log_deprecation_warning(_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING)
+
+        # These attributes should only be set through `BaseTrainer.restore`
+        self._restore_path = None
+        self._restore_storage_filesystem = None
+
+        self._validate_attributes()
+
+        usage_lib.record_library_usage("train")
+        air_usage.tag_air_trainer(self)
+
+    @classmethod
+    @Deprecated(message=_TRAINER_RESTORE_DEPRECATION_WARNING)
+    def restore(
+        cls: Type["BaseTrainer"],
+        path: Union[str, os.PathLike],
+        storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        scaling_config: Optional[ScalingConfig] = None,
+        **kwargs,
+    ) -> "BaseTrainer":
+        """Restores a Train experiment from a previously interrupted/failed run.
+
+        Restore should be used for experiment-level fault tolerance in the event
+        that the head node crashes (e.g., OOM or some other runtime error) or the
+        entire cluster goes down (e.g., network error affecting all nodes).
+
+        A run that has already completed successfully will not be resumed from this API.
+        To continue training from a successful run, launch a new run with the
+        ``<Framework>Trainer(resume_from_checkpoint)`` API instead, passing in a
+        checkpoint from the previous run to start with.
+
+        .. warning::
+
+            The ``path`` must point to a **trusted** experiment directory.
+            Restoring from an untrusted path executes arbitrary Python code
+            (the experiment state uses pickle serialization). Never restore
+            from a path that other parties can write to.
+
+        .. note::
+
+            Restoring an experiment from a path that's pointing to a *different*
+            location than the original experiment path is supported. However, Ray Train
+            assumes that the full experiment directory is available
+            (including checkpoints) so that it's possible to resume trials from their
+            latest state.
+
+            For example, if the original experiment path was run locally, then the
+            results are uploaded to cloud storage, Ray Train expects the full contents
+            to be available in cloud storage if attempting to resume
+            via ``<Framework>Trainer.restore("s3://...")``. The restored run will
+            continue writing results to the same cloud storage location.
+
+        The following example can be paired with implementing job retry using
+        :ref:`Ray Jobs <jobs-overview>` to produce a Train experiment that will
+        attempt to resume on both experiment-level and trial-level failures:
+
+        .. testcode::
+            :skipif: True
+
+            import os
+            import ray
+            from ray import train
+            from ray.train.trainer import BaseTrainer
+
+            experiment_name = "unique_experiment_name"
+            storage_path = os.path.expanduser("~/ray_results")
+            experiment_dir = os.path.join(storage_path, experiment_name)
+
+            # Define some dummy inputs for demonstration purposes
+            datasets = {"train": ray.data.from_items([{"a": i} for i in range(10)])}
+
+            class CustomTrainer(BaseTrainer):
+                def training_loop(self):
+                    pass
+
+            if CustomTrainer.can_restore(experiment_dir):
+                trainer = CustomTrainer.restore(
+                    experiment_dir, datasets=datasets
+                )
+            else:
+                trainer = CustomTrainer(
+                    datasets=datasets,
+                    run_config=train.RunConfig(
+                        name=experiment_name,
+                        storage_path=storage_path,
+                        # Tip: You can also enable retries on failure for
+                        # worker-level fault tolerance
+                        failure_config=train.FailureConfig(max_failures=3),
+                    ),
+                )
+
+            result = trainer.fit()
+
+        Args:
+            path: The path to the experiment directory of the training run to restore.
+                This can be a local path or a remote URI if the experiment was
+                uploaded to the cloud.
+            storage_filesystem: Custom ``pyarrow.fs.FileSystem``
+                corresponding to the ``path``. This may be necessary if the original
+                experiment passed in a custom filesystem.
+            datasets: Re-specified datasets used in the original training run.
+                This must include all the datasets that were passed in the
+                original trainer constructor.
+            scaling_config: Optionally re-specified scaling config. This can be
+                modified to be different from the original spec.
+            **kwargs: Other optionally re-specified arguments, passed in by subclasses.
+
+        Raises:
+            ValueError: If all datasets were not re-supplied on restore.
+
+        Returns:
+            BaseTrainer: A restored instance of the class that is calling this method.
+        """
+        if _v2_migration_warnings_enabled():
+            _log_deprecation_warning(_TRAINER_RESTORE_DEPRECATION_WARNING)
+
+        if not cls.can_restore(path, storage_filesystem):
+            raise ValueError(
+                f"Invalid restore path: {path}. Make sure that this path exists and "
+                "is the experiment directory that results from a call to "
+                "`trainer.fit()`."
+            )
+        fs, fs_path = get_fs_and_path(path, storage_filesystem)
+        trainer_pkl_path = Path(fs_path, _TRAINER_PKL).as_posix()
+        with fs.open_input_file(trainer_pkl_path) as f:
+            trainer_cls, param_dict = pickle.loads(f.readall())
+
+        if trainer_cls is not cls:
+            warnings.warn(
+                f"Invalid trainer type. You are attempting to restore a trainer of type"
+                f" {trainer_cls} with `{cls.__name__}.restore`, "
+                "which will most likely fail. "
+                f"Use `{trainer_cls.__name__}.restore` instead."
+            )
+
+        original_datasets = param_dict.pop("datasets", {})
+        if original_datasets and not datasets:
+            raise ValueError(
+                "The following datasets need to be provided again on restore: "
+                f"{list(original_datasets.keys())}\n"
+                f"Use {cls.__name__}.restore(..., datasets=datasets) "
+                "with the datasets that were provided to the original trainer."
+            )
+        datasets = datasets or {}
+        if set(original_datasets) != set(datasets):
+            raise ValueError(
+                "The provided datasets don't match the original dataset keys.\n"
+                f"  Expected datasets for the keys: {list(original_datasets.keys())}\n"
+                f"  Actual datasets provided: {list(datasets.keys())}"
+            )
+        param_dict["datasets"] = datasets
+
+        if scaling_config:
+            param_dict["scaling_config"] = scaling_config
+
+        for param_name, val in kwargs.items():
+            # Overwrite the old value if something is passed into restore
+            if val is not None:
+                param_dict[param_name] = val
+
+        try:
+            trainer = cls(**param_dict)
+        except Exception as e:
+            raise ValueError(
+                "Trainer restoration failed (see above for the stack trace). "
+                "Make sure that you use the right trainer class to restore: "
+                f"`{cls.__name__}.restore`\n"
+            ) from e
+        trainer._restore_path = path
+        trainer._restore_storage_filesystem = storage_filesystem
+        return trainer
+
+    @classmethod
+    @Deprecated(
+        message=_TRAINER_RESTORE_DEPRECATION_WARNING,
+        warning=_v2_migration_warnings_enabled(),
+    )
+    def can_restore(
+        cls: Type["BaseTrainer"],
+        path: Union[str, os.PathLike],
+        storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
+    ) -> bool:
+        """Checks whether a given directory contains a restorable Train experiment.
+
+        Args:
+            path: The path to the experiment directory of the Train experiment.
+                This can be either a local directory (e.g., ~/ray_results/exp_name)
+                or a remote URI (e.g., s3://bucket/exp_name).
+            storage_filesystem: Custom ``pyarrow.fs.FileSystem`` to use. If not
+                provided, the filesystem is auto-resolved from ``path``.
+
+        Returns:
+            bool: Whether this path exists and contains the trainer state to resume from
+        """
+        if _v2_migration_warnings_enabled():
+            _log_deprecation_warning(_TRAINER_RESTORE_DEPRECATION_WARNING)
+
+        fs, fs_path = get_fs_and_path(path, storage_filesystem)
+        trainer_pkl_path = Path(fs_path, _TRAINER_PKL).as_posix()
+        return _exists_at_fs_path(fs, trainer_pkl_path)
+
+    def __repr__(self):
+        # A dictionary that maps parameters to their default values.
+        default_values: Dict[str, Any] = {
+            "scaling_config": ScalingConfig(),
+            "run_config": RunConfig(),
+            "datasets": {},
+            "starting_checkpoint": None,
+        }
+
+        non_default_arguments = []
+        for parameter, default_value in default_values.items():
+            value = getattr(self, parameter)
+            if value != default_value:
+                # 'Dataset.__repr__' returns a table rather than a regular Python object
+                # representation. So, we need to special case the 'datasets' parameter.
+                if parameter == "datasets":
+                    value_repr = format_datasets_for_repr(value)
+                else:
+                    value_repr = repr(value)
+
+                non_default_arguments.append(f"{parameter}={value_repr}")
+
+        if non_default_arguments:
+            return f"<{self.__class__.__name__} {' '.join(non_default_arguments)}>"
+
+        return f"<{self.__class__.__name__}>"
+
+    def __new__(cls, *args, **kwargs):
+        # Store the init args as attributes so this can be merged with Tune hparams.
+        trainer = super(BaseTrainer, cls).__new__(cls)
+        parameters = inspect.signature(cls.__init__).parameters
+        parameters = list(parameters.keys())
+        # Remove self.
+        parameters = parameters[1:]
+        arg_dict = dict(zip(parameters, args))
+        trainer._param_dict = {**arg_dict, **kwargs}
+        return trainer
+
+    def _validate_attributes(self):
+        """Called on __init()__ to validate trainer attributes."""
+        # Run config
+        if not isinstance(self.run_config, RunConfig):
+            raise ValueError(
+                f"`run_config` should be an instance of `ray.train.RunConfig`, "
+                f"found {type(self.run_config)} with value `{self.run_config}`."
+            )
+        # Scaling config
+        if not isinstance(self.scaling_config, ScalingConfig):
+            raise ValueError(
+                "`scaling_config` should be an instance of `ScalingConfig`, "
+                f"found {type(self.scaling_config)} with value `{self.scaling_config}`."
+            )
+        # Datasets
+        if not isinstance(self.datasets, dict):
+            raise ValueError(
+                f"`datasets` should be a dict mapping from a string to "
+                f"`ray.data.Dataset` objects, "
+                f"found {type(self.datasets)} with value `{self.datasets}`."
+            )
+        else:
+            for key, dataset in self.datasets.items():
+                if not isinstance(dataset, ray.data.Dataset) and not callable(dataset):
+                    raise ValueError(
+                        f"The Dataset under '{key}' key is not a "
+                        "`ray.data.Dataset`. "
+                        f"Received {dataset} instead."
+                    )
+        # Metadata.
+        self.metadata = self.metadata or {}
+        if not isinstance(self.metadata, dict):
+            raise TypeError(
+                f"The provided metadata must be a dict, was {type(self.metadata)}."
+            )
+        try:
+            self.metadata = json.loads(json.dumps(self.metadata))
+        except Exception as e:
+            raise ValueError(
+                "The provided metadata must be JSON-serializable: "
+                f"{self.metadata}: {e}"
+            )
+
+        if self.starting_checkpoint is not None and not isinstance(
+            self.starting_checkpoint, Checkpoint
+        ):
+            raise ValueError(
+                f"`resume_from_checkpoint` should be an instance of "
+                f"`ray.train.Checkpoint`, found {type(self.starting_checkpoint)} "
+                f"with value `{self.starting_checkpoint}`."
+            )
+
+        self._log_v2_deprecation_warnings()
+
+    def _log_v2_deprecation_warnings(self):
+        """Logs deprecation warnings for v2 migration.
+
+        Log them here in the Ray Train case rather than in the configuration
+        constructors to avoid logging incorrect deprecation warnings when
+        `ray.train.RunConfig` is passed to Ray Tune.
+        """
+        from ray.train.v2._internal.constants import V2_ENABLED_ENV_VAR, is_v2_enabled
+
+        if is_v2_enabled():
+            raise DeprecationWarning(
+                f"Detected use of a deprecated Trainer import from `{self.__class__.__module__}`. "
+                "This Trainer class is not compatible with Ray Train V2.\n"
+                "To fix this:\n"
+                "  - Update to use the new import path. For example, "
+                "`from ray.train.torch.torch_trainer import TorchTrainer` -> "
+                "`from ray.train.torch import TorchTrainer`\n"
+                f"  - Or, explicitly disable V2 by setting: {V2_ENABLED_ENV_VAR}=0\n"
+                "See this issue for more context: "
+                "https://github.com/ray-project/ray/issues/49454"
+            )
+
+        if not _v2_migration_warnings_enabled():
+            return
+
+        from ray.train.v2._internal.migration_utils import (
+            CALLBACKS_DEPRECATION_MESSAGE,
+            FAIL_FAST_DEPRECATION_MESSAGE,
+            LOG_TO_FILE_DEPRECATION_MESSAGE,
+            PROGRESS_REPORTER_DEPRECATION_MESSAGE,
+            STOP_DEPRECATION_MESSAGE,
+            SYNC_CONFIG_DEPRECATION_MESSAGE,
+            TRAINER_RESOURCES_DEPRECATION_MESSAGE,
+            VERBOSE_DEPRECATION_MESSAGE,
+        )
+
+        # ScalingConfig deprecations
+        if self.scaling_config.trainer_resources is not None:
+            _log_deprecation_warning(TRAINER_RESOURCES_DEPRECATION_MESSAGE)
+
+        # FailureConfig deprecations
+        if self.run_config.failure_config.fail_fast:
+            _log_deprecation_warning(FAIL_FAST_DEPRECATION_MESSAGE)
+
+        # RunConfig deprecations
+        # NOTE: _verbose is the original verbose value passed by the user
+        if self.run_config._verbose is not None:
+            _log_deprecation_warning(VERBOSE_DEPRECATION_MESSAGE)
+
+        if self.run_config.log_to_file:
+            _log_deprecation_warning(LOG_TO_FILE_DEPRECATION_MESSAGE)
+
+        if self.run_config.stop is not None:
+            _log_deprecation_warning(STOP_DEPRECATION_MESSAGE)
+
+        if self.run_config.callbacks is not None:
+            _log_deprecation_warning(CALLBACKS_DEPRECATION_MESSAGE)
+
+        if self.run_config.progress_reporter is not None:
+            _log_deprecation_warning(PROGRESS_REPORTER_DEPRECATION_MESSAGE)
+
+        if self.run_config.sync_config != ray.train.SyncConfig():
+            _log_deprecation_warning(SYNC_CONFIG_DEPRECATION_MESSAGE)
+
+    @classmethod
+    def _validate_scaling_config(cls, scaling_config: ScalingConfig) -> ScalingConfig:
+        """Returns scaling config dataclass after validating updated keys."""
+        ensure_only_allowed_dataclass_keys_updated(
+            dataclass=scaling_config,
+            allowed_keys=cls._scaling_config_allowed_keys,
+        )
+        return scaling_config
+
+    def setup(self) -> None:
+        """Called during fit() to perform initial setup on the Trainer.
+
+        .. note:: This method is run on a remote process.
+
+        This method will not be called on the driver, so any expensive setup
+        operations should be placed here and not in ``__init__``.
+
+        This method is called prior to ``preprocess_datasets`` and
+        ``training_loop``.
+        """
+        pass
+
+    def preprocess_datasets(self) -> None:
+        """Deprecated."""
+        raise DeprecationWarning(
+            "`preprocess_datasets` is no longer used, since preprocessors "
+            f"are no longer accepted by Trainers.\n{PREPROCESSOR_DEPRECATION_MESSAGE}"
+        )
+
+    @abc.abstractmethod
+    def training_loop(self) -> None:
+        """Loop called by fit() to run training and report results to Tune.
+
+        .. note:: This method runs on a remote process.
+
+        ``self.datasets`` have already been evaluated if they were wrapped in a factory.
+
+        You can use the :ref:`Ray Train utilities <train-loop-api>`
+        (:func:`train.report() <ray.train.report>` and
+        :func:`train.get_checkpoint() <ray.train.get_checkpoint>`) inside
+        this training loop.
+
+        Example:
+
+        .. testcode::
+
+            from ray.train.trainer import BaseTrainer
+            from ray import train
+
+            class MyTrainer(BaseTrainer):
+                def training_loop(self):
+                    for epoch_idx in range(5):
+                        ...
+                        train.report({"epoch": epoch_idx})
+
+        """
+        raise NotImplementedError
+
+    @PublicAPI(stability="beta")
+    def fit(self) -> Result:
+        """Runs training.
+
+        Returns:
+            A Result object containing the training result.
+
+        Raises:
+            ray.train.base_trainer.TrainingFailedError: If any failures during the execution
+                of ``self.as_trainable()``, or during the Tune execution loop.
+        """
+        from ray.tune import ResumeConfig, TuneError
+        from ray.tune.tuner import Tuner
+
+        trainable = self.as_trainable()
+        param_space = self._extract_fields_for_tuner_param_space()
+
+        self.run_config.name = (
+            self.run_config.name or StorageContext.get_experiment_dir_name(trainable)
+        )
+        # The storage context here is only used to access the resolved
+        # storage fs and experiment path, in order to avoid duplicating that logic.
+        # This is NOT the storage context object that gets passed to remote workers.
+        storage = StorageContext(
+            storage_path=self.run_config.storage_path,
+            experiment_dir_name=self.run_config.name,
+            storage_filesystem=self.run_config.storage_filesystem,
+        )
+
+        if self._restore_path:
+            tuner = Tuner.restore(
+                path=self._restore_path,
+                trainable=trainable,
+                param_space=param_space,
+                _resume_config=ResumeConfig(
+                    finished=ResumeConfig.ResumeType.RESUME,
+                    unfinished=ResumeConfig.ResumeType.RESUME,
+                    errored=ResumeConfig.ResumeType.RESUME,
+                ),
+                storage_filesystem=self._restore_storage_filesystem,
+            )
+        else:
+            tuner = Tuner(
+                trainable=trainable,
+                param_space=param_space,
+                run_config=self.run_config,
+                _entrypoint=AirEntrypoint.TRAINER,
+            )
+
+        self._save(storage.storage_filesystem, storage.experiment_fs_path)
+
+        restore_msg = TrainingFailedError._RESTORE_MSG.format(
+            trainer_cls_name=self.__class__.__name__,
+            path=str(storage.experiment_fs_path),
+        )
+
+        try:
+            result_grid = tuner.fit()
+        except TuneError as e:
+            # Catch any `TuneError`s raised by the `Tuner.fit` call.
+            # Unwrap the `TuneError` if needed.
+            parent_error = e.__cause__ or e
+
+            # Raise it to the user as a `TrainingFailedError` with a message to restore.
+            raise TrainingFailedError(restore_msg) from parent_error
+        # Other exceptions get passed through directly (ex: on `fail_fast='raise'`)
+
+        assert len(result_grid) == 1
+        result = result_grid[0]
+        if result.error:
+            # Raise trainable errors to the user with a message to restore
+            # or configure `FailureConfig` in a new run.
+            raise TrainingFailedError(
+                "\n".join([restore_msg, TrainingFailedError._FAILURE_CONFIG_MSG])
+            ) from result.error
+        return result
+
+    def _save(self, fs: pyarrow.fs.FileSystem, experiment_path: str):
+        """Saves the current trainer's class along with the `param_dict` of
+        parameters passed to this trainer's constructor.
+
+        This is used to recreate the trainer on restore.
+        Unless a parameter is re-specified during restoration (only a subset
+        of parameters can be passed in again), that parameter will be loaded
+        from the saved copy.
+
+        Datasets should not be saved as part of the state. Instead, we save the
+        keys and replace the dataset values with dummy functions that will
+        raise an error if invoked. The error only serves as a guardrail for
+        misuse (e.g., manually unpickling and constructing the Trainer again)
+        and is not typically surfaced, since datasets must be re-specified
+        upon restoration.
+        """
+        param_dict = self._param_dict.copy()
+        datasets = param_dict.pop("datasets", {})
+
+        def raise_fn():
+            raise RuntimeError
+
+        if datasets:
+            param_dict["datasets"] = dict.fromkeys(datasets, raise_fn)
+
+        cls_and_param_dict = (self.__class__, param_dict)
+
+        fs.create_dir(experiment_path)
+        with fs.open_output_stream(Path(experiment_path, _TRAINER_PKL).as_posix()) as f:
+            f.write(pickle.dumps(cls_and_param_dict))
+
+    def _extract_fields_for_tuner_param_space(self) -> Dict:
+        """Extracts fields to be included in `Tuner.param_space`.
+
+        This is needed to leverage the full logging/integration offerings from Tune.
+        For example, `param_space` is logged automatically to wandb integration.
+
+        Currently only done for `train_loop_config`.
+
+        Returns:
+            A dictionary that should be passed to Tuner.param_space.
+        """
+        result = {}
+        for key in self._fields_for_tuner_param_space:
+            if key in self._param_dict.keys():
+                result[key] = copy.deepcopy(self._param_dict[key])
+        return result
+
+    def _generate_trainable_cls(self) -> Type["Trainable"]:
+        """Generates the base Trainable class.
+
+        Returns:
+            A Trainable class to use for training.
+        """
+
+        from ray.tune.execution.placement_groups import PlacementGroupFactory
+        from ray.tune.trainable import wrap_function
+
+        trainer_cls = self.__class__
+        scaling_config = self.scaling_config
+        metadata = self.metadata
+
+        train_coordinator_fn = partial(
+            _train_coordinator_fn, trainer_cls=trainer_cls, metadata=metadata
+        )
+        # Change the name of the training function to match the name of the Trainer
+        # class. This will mean the Tune trial name will match the name of Trainer on
+        # stdout messages and the results directory.
+        train_coordinator_fn.__name__ = trainer_cls.__name__
+
+        trainable_cls = wrap_function(train_coordinator_fn)
+        has_base_dataset = bool(self.datasets)
+        if has_base_dataset:
+            from ray.data.context import DataContext
+
+            dataset_context = DataContext.get_current()
+        else:
+            dataset_context = None
+
+        class TrainTrainable(trainable_cls):
+            """Adds default resources to the Trainable."""
+
+            _handles_checkpoint_freq = trainer_cls._handles_checkpoint_freq
+            _handles_checkpoint_at_end = trainer_cls._handles_checkpoint_at_end
+
+            @classmethod
+            def has_base_dataset(cls) -> bool:
+                """Whether a dataset is provided through the Trainer."""
+                return has_base_dataset
+
+            @classmethod
+            def base_scaling_config(cls) -> ScalingConfig:
+                """Returns the unchanged scaling config provided through the Trainer."""
+                return scaling_config
+
+            def setup(self, config, **kwargs):
+                base_config = dict(kwargs)
+                # Merge Tuner param space hyperparameters in `config` into the
+                # base config passed to the Trainer constructor, which is `base_config`.
+                # `base_config` is pulled from the object store from the usage of
+                # tune.with_parameters in `BaseTrainer.as_trainable`.
+
+                # run_config is not a tunable hyperparameter so it does not need to be
+                # merged.
+                run_config = base_config.pop("run_config", None)
+                self._merged_config = deep_update(
+                    base_config, self.config, new_keys_allowed=True
+                )
+                self._merged_config["run_config"] = run_config
+                merged_scaling_config = self._merged_config.get(
+                    "scaling_config", ScalingConfig()
+                )
+                if isinstance(merged_scaling_config, dict):
+                    merged_scaling_config = ScalingConfig(**merged_scaling_config)
+                self._merged_config[
+                    "scaling_config"
+                ] = self._reconcile_scaling_config_with_trial_resources(
+                    merged_scaling_config
+                )
+                if self.has_base_dataset():
+                    # Set the DataContext on the Trainer actor to the DataContext
+                    # specified on the driver.
+                    DataContext._set_current(dataset_context)
+                super(TrainTrainable, self).setup(config)
+
+            def _reconcile_scaling_config_with_trial_resources(
+                self, scaling_config: ScalingConfig
+            ) -> ScalingConfig:
+                """
+                ResourceChangingScheduler workaround.
+
+                Ensures that the scaling config matches trial resources.
+
+                This should be replaced with RCS returning a ScalingConfig
+                in the future.
+                """
+
+                trial_resources = self.trial_resources
+                # This will be false if the resources are default
+                if not isinstance(trial_resources, PlacementGroupFactory):
+                    return scaling_config
+
+                # Ignore ResourceChangingScheduler workaround when resource bundles
+                # are unchanged
+                if self.trial_resources == scaling_config.as_placement_group_factory():
+                    return scaling_config
+
+                trainer_cls._validate_scaling_config(scaling_config)
+
+                return ScalingConfig.from_placement_group_factory(trial_resources)
+
+            def _trainable_func(self, config):
+                # We ignore the config passed by Tune and instead use the merged
+                # config which includes the initial Trainer args.
+                super()._trainable_func(self._merged_config)
+
+            @classmethod
+            def default_resource_request(cls, config):
+                # `config["scaling_config"] is a dataclass when passed via the
+                # `scaling_config` argument in `Trainer` and is a dict when passed
+                # via the `scaling_config` key of `param_spec`.
+
+                # Conversion logic must be duplicated in `TrainTrainable.__init__`
+                # because this is a class method.
+                updated_scaling_config = config.get("scaling_config", scaling_config)
+                if isinstance(updated_scaling_config, dict):
+                    updated_scaling_config = ScalingConfig(**updated_scaling_config)
+                validated_scaling_config = trainer_cls._validate_scaling_config(
+                    updated_scaling_config
+                )
+                return validated_scaling_config.as_placement_group_factory()
+
+        return TrainTrainable
+
+    def as_trainable(self) -> Type["Trainable"]:
+        """Converts self to a ``tune.Trainable`` class."""
+        from ray import tune
+
+        base_config = self._param_dict
+        trainable_cls = self._generate_trainable_cls()
+
+        # Wrap with `tune.with_parameters` to handle very large values in base_config
+        return tune.with_parameters(trainable_cls, **base_config)
+
+
+@DeveloperAPI
+def format_datasets_for_repr(datasets: Optional[Dict[str, GenDataset]]) -> str:
+    """Format datasets for BaseTrainer repr using plan strings.
+
+    The Dataset.__repr__ returns a table rather than a conventional Python object
+    reprentation. To ensure the BaseTrainer representation still looks reasonable, we
+    need to special-case datasets.
+    """
+    from ray.data import Dataset
+    from ray.data._internal.dataset_repr import build_dataset_summary_repr
+
+    assert datasets is not None, "Expected caller to pass in non-None argument"
+
+    formatted = {}
+    for key, dataset in datasets.items():
+        if isinstance(dataset, Dataset):
+            formatted[key] = build_dataset_summary_repr(dataset)
+        else:
+            formatted[key] = dataset
+
+    return "{" + ", ".join(f"'{key}': {formatted[key]}" for key in datasets) + "}"

python/ray/train/collective/__init__.py

@@ -0,0 +1,19 @@
+from ray.train.v2._internal.constants import is_v2_enabled
+
+if is_v2_enabled():
+    from ray.train.collective.collectives import barrier, broadcast_from_rank_zero
+
+    __all__ = [
+        "broadcast_from_rank_zero",
+        "barrier",
+    ]
+
+    broadcast_from_rank_zero.__module__ = "ray.train.collective"
+    barrier.__module__ = "ray.train.collective"
+else:
+    raise ImportError(
+        "`ray.train.collective` is only available in Ray Train v2. "
+        "To enable it, please set `RAY_TRAIN_V2_ENABLED=1`."
+    )
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/collective/collectives.py

@@ -0,0 +1,82 @@
+import logging
+from typing import Optional, TypeVar
+
+from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
+from ray.train.v2._internal.util import requires_train_worker
+from ray.util.annotations import PublicAPI
+
+T = TypeVar("T", bound=Optional[object])
+
+
+logger = logging.getLogger(__file__)
+
+
+@PublicAPI(stability="alpha")
+@requires_train_worker()
+def broadcast_from_rank_zero(data: T) -> T:
+    """Broadcast small (<1kb) data from the rank 0 worker to all other workers.
+
+    Serves as a barrier, meaning that all workers must call this method before
+    the training function can continue.
+
+    Example:
+
+        .. testcode:
+
+            from ray.train import get_context
+            from ray.train.collective import broadcast_from_rank_zero
+            from ray.train.torch import TorchTrainer
+
+            def train_func():
+                ...
+                if get_context().get_world_rank() == 0:
+                    data = {"some_key": "some_value"}
+                else:
+                    data = None
+                data = broadcast_from_rank_zero(data)
+                ...
+
+            trainer = TorchTrainer(train_func)
+            trainer.fit()
+
+    Args:
+        data: The small (1kb) data to broadcast from the rank 0 worker to all
+            other workers.
+
+    Returns:
+        The data broadcasted from the rank 0 worker.
+
+    Raises:
+        ValueError: If the data is too big.
+        pickle.PicklingError: If the data is not pickleable.
+        TypeError: If the data is not pickleable.
+    """
+    return get_train_fn_utils().broadcast_from_rank_zero(data)
+
+
+@PublicAPI(stability="alpha")
+@requires_train_worker()
+def barrier() -> None:
+    """Create a barrier across all workers.
+
+    All workers must call this method before the training function can continue.
+
+    Example:
+
+        .. testcode:
+
+            from ray.train import get_context
+            from ray.train.collective import barrier
+            from ray.train.torch import TorchTrainer
+
+            def train_func():
+                ...
+                print(f"Rank {get_context().get_world_rank()} is waiting at the barrier.")
+                barrier()
+                print(f"Rank {get_context().get_world_rank()} has passed the barrier.")
+                ...
+
+            trainer = TorchTrainer(train_func)
+            trainer.fit()
+    """
+    return get_train_fn_utils().barrier()

python/ray/train/constants.py

@@ -0,0 +1,158 @@
+from pathlib import Path
+from typing import Any
+
+import ray
+from ray._private.ray_constants import env_bool
+from ray.air.constants import (  # noqa: F401
+    COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV,
+    EVALUATION_DATASET_KEY,
+    MODEL_KEY,
+    PREPROCESSOR_KEY,
+    TRAIN_DATASET_KEY,
+)
+
+
+def _get_ray_train_session_dir() -> str:
+    assert ray.is_initialized(), "Ray must be initialized to get the session dir."
+    return Path(
+        ray._private.worker._global_node.get_session_dir_path(), "artifacts"
+    ).as_posix()
+
+
+DEFAULT_STORAGE_PATH = Path("~/ray_results").expanduser().as_posix()
+
+# Autofilled ray.train.report() metrics. Keys should be consistent with Tune.
+CHECKPOINT_DIR_NAME = "checkpoint_dir_name"
+TIME_TOTAL_S = "_time_total_s"
+WORKER_HOSTNAME = "_hostname"
+WORKER_NODE_IP = "_node_ip"
+WORKER_PID = "_pid"
+
+# Will not be reported unless ENABLE_DETAILED_AUTOFILLED_METRICS_ENV
+# env var is not 0
+DETAILED_AUTOFILLED_KEYS = {WORKER_HOSTNAME, WORKER_NODE_IP, WORKER_PID, TIME_TOTAL_S}
+
+# Default filename for JSON logger
+RESULT_FILE_JSON = "results.json"
+
+# The name of the subdirectory inside the trainer run_dir to store checkpoints.
+TRAIN_CHECKPOINT_SUBDIR = "checkpoints"
+
+# The key to use to specify the checkpoint id for Tune.
+# This needs to be added to the checkpoint dictionary so if the Tune trial
+# is restarted, the checkpoint_id can continue to increment.
+TUNE_CHECKPOINT_ID = "_current_checkpoint_id"
+
+# Deprecated configs can use this value to detect if the user has set it.
+# This has type Any to allow it to be assigned to any annotated parameter
+# without causing type errors.
+_DEPRECATED_VALUE: Any = "DEPRECATED"
+
+
+# ==================================================
+#               Train V2 constants
+# ==================================================
+
+# Set this to 1 to enable deprecation warnings for V2 migration.
+ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR = "RAY_TRAIN_ENABLE_V2_MIGRATION_WARNINGS"
+
+
+V2_MIGRATION_GUIDE_MESSAGE = (
+    "See this issue for more context and migration options: "
+    "https://github.com/ray-project/ray/issues/49454. "
+    "Disable these warnings by setting the environment variable: "
+    f"{ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR}=0"
+)
+
+
+def _v2_migration_warnings_enabled() -> bool:
+    return env_bool(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR, True)
+
+
+# ==================================================
+#               Environment Variables
+# ==================================================
+
+ENABLE_DETAILED_AUTOFILLED_METRICS_ENV = (
+    "TRAIN_RESULT_ENABLE_DETAILED_AUTOFILLED_METRICS"
+)
+
+# Integer value which if set will override the value of
+# Backend.share_cuda_visible_devices. 1 for True, 0 for False.
+ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV = "TRAIN_ENABLE_SHARE_CUDA_VISIBLE_DEVICES"
+
+# Integer value which if set will not share HIP accelerator visible devices
+# across workers. 1 for True (default), 0 for False.
+ENABLE_SHARE_HIP_VISIBLE_DEVICES_ENV = "TRAIN_ENABLE_SHARE_HIP_VISIBLE_DEVICES"
+
+# Integer value which if set will not share neuron-core accelerator visible cores
+# across workers. 1 for True (default), 0 for False.
+ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV = (
+    "TRAIN_ENABLE_SHARE_NEURON_CORES_ACCELERATOR"
+)
+
+# Integer value which if set will not share npu visible devices
+# across workers. 1 for True (default), 0 for False.
+ENABLE_SHARE_NPU_RT_VISIBLE_DEVICES_ENV = "TRAIN_ENABLE_SHARE_ASCEND_RT_VISIBLE_DEVICES"
+
+# Integer value which indicates the number of seconds to wait when creating
+# the worker placement group before timing out.
+TRAIN_PLACEMENT_GROUP_TIMEOUT_S_ENV = "TRAIN_PLACEMENT_GROUP_TIMEOUT_S"
+
+# Integer value which if set will change the placement group strategy from
+# PACK to SPREAD. 1 for True, 0 for False.
+TRAIN_ENABLE_WORKER_SPREAD_ENV = "TRAIN_ENABLE_WORKER_SPREAD"
+
+# Set this to 0 to disable changing the working directory of each Tune Trainable
+# or Train worker to the trial directory. Defaults to 1.
+RAY_CHDIR_TO_TRIAL_DIR = "RAY_CHDIR_TO_TRIAL_DIR"
+
+# Set this to 1 to count preemption errors toward `FailureConfig(max_failures)`.
+# Defaults to 0, which always retries on node preemption failures.
+RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE = "RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE"
+
+# Set this to 1 to start a StateActor and collect information Train Runs
+# Defaults to 0
+RAY_TRAIN_ENABLE_STATE_TRACKING = "RAY_TRAIN_ENABLE_STATE_TRACKING"
+
+# Set this to 1 to only store the checkpoint score attribute with the Checkpoint
+# in the CheckpointManager. The Result will only have the checkpoint score attribute
+# but files written to disk like result.json will still have all the metrics.
+# Defaults to 0.
+# TODO: this is a temporary solution to avoid CheckpointManager OOM.
+# See https://github.com/ray-project/ray/pull/54642#issue-3234029360 for more details.
+TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE = (
+    "TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE"
+)
+
+# Seconds to wait for torch process group to shut down.
+# Shutting down a healthy torch process group, which we may want to do for reasons
+# like restarting a group of workers if an async checkpoint upload fails, can hang.
+# This is a workaround until we figure out how to avoid this hang.
+TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S = "TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S"
+DEFAULT_TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S = 30
+
+# Seconds to wait for JAX distributed shutdown.
+JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S = "JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S"
+DEFAULT_JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S = 30
+
+# NOTE: When adding a new environment variable, please track it in this list.
+TRAIN_ENV_VARS = {
+    ENABLE_DETAILED_AUTOFILLED_METRICS_ENV,
+    ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
+    ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
+    TRAIN_PLACEMENT_GROUP_TIMEOUT_S_ENV,
+    TRAIN_ENABLE_WORKER_SPREAD_ENV,
+    RAY_CHDIR_TO_TRIAL_DIR,
+    RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE,
+    RAY_TRAIN_ENABLE_STATE_TRACKING,
+    TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE,
+    TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
+    JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
+}
+
+# Key for AIR Checkpoint metadata in TrainingResult metadata
+CHECKPOINT_METADATA_KEY = "checkpoint_metadata"
+
+# Key for AIR Checkpoint world rank in TrainingResult metadata
+CHECKPOINT_RANK_KEY = "checkpoint_rank"

python/ray/train/context.py

@@ -0,0 +1,144 @@
+import threading
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from ray.train._internal import session
+from ray.train._internal.storage import StorageContext
+from ray.train.constants import (
+    V2_MIGRATION_GUIDE_MESSAGE,
+    _v2_migration_warnings_enabled,
+)
+from ray.train.utils import _copy_doc, _log_deprecation_warning
+from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.execution.placement_groups import PlacementGroupFactory
+
+
+# The context singleton on this process.
+_default_context: "Optional[TrainContext]" = None
+_context_lock = threading.Lock()
+
+
+_GET_METADATA_DEPRECATION_MESSAGE = (
+    "`get_metadata` was an experimental API that accessed the metadata passed "
+    "to `<Framework>Trainer(metadata=...)`. This API can be replaced by passing "
+    "the metadata directly to the training function (e.g., via `train_loop_config`). "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE = (
+    "`{}` is deprecated because the concept of a `Trial` will "
+    "soon be removed in Ray Train."
+    "Ray Train will no longer assume that it's running within a Ray Tune `Trial` "
+    "in the future. "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+
+@PublicAPI(stability="stable")
+class TrainContext:
+    """Context containing metadata that can be accessed within Ray Train workers."""
+
+    @_copy_doc(session.get_experiment_name)
+    def get_experiment_name(self) -> str:
+        return session.get_experiment_name()
+
+    @_copy_doc(session.get_world_size)
+    def get_world_size(self) -> int:
+        return session.get_world_size()
+
+    @_copy_doc(session.get_world_rank)
+    def get_world_rank(self) -> int:
+        return session.get_world_rank()
+
+    @_copy_doc(session.get_local_rank)
+    def get_local_rank(self) -> int:
+        return session.get_local_rank()
+
+    @_copy_doc(session.get_local_world_size)
+    def get_local_world_size(self) -> int:
+        return session.get_local_world_size()
+
+    @_copy_doc(session.get_node_rank)
+    def get_node_rank(self) -> int:
+        return session.get_node_rank()
+
+    @DeveloperAPI
+    @_copy_doc(session.get_storage)
+    def get_storage(self) -> StorageContext:
+        return session.get_storage()
+
+    # Deprecated APIs
+
+    @Deprecated(
+        message=_GET_METADATA_DEPRECATION_MESSAGE,
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(session.get_metadata)
+    def get_metadata(self) -> Dict[str, Any]:
+        return session.get_metadata()
+
+    @Deprecated(
+        message=_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_name"),
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(session.get_trial_name)
+    def get_trial_name(self) -> str:
+        return session.get_trial_name()
+
+    @Deprecated(
+        message=_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_id"),
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(session.get_trial_id)
+    def get_trial_id(self) -> str:
+        return session.get_trial_id()
+
+    @Deprecated(
+        message=_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format(
+            "get_trial_resources"
+        ),
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(session.get_trial_resources)
+    def get_trial_resources(self) -> "PlacementGroupFactory":
+        return session.get_trial_resources()
+
+    @Deprecated(
+        message=_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_dir"),
+        warning=_v2_migration_warnings_enabled(),
+    )
+    @_copy_doc(session.get_trial_dir)
+    def get_trial_dir(self) -> str:
+        return session.get_trial_dir()
+
+
+@PublicAPI(stability="stable")
+def get_context() -> TrainContext:
+    """Get or create a singleton training context.
+
+    The context is only available within a function passed to Ray Train.
+
+    See the :class:`~ray.train.TrainContext` API reference to see available methods.
+    """
+    from ray.tune.trainable.trainable_fn_utils import _in_tune_session
+
+    # If we are running in a Tune function, switch to Tune context.
+    if _in_tune_session():
+        from ray.tune import get_context as get_tune_context
+
+        if _v2_migration_warnings_enabled():
+            _log_deprecation_warning(
+                "`ray.train.get_context()` should be switched to "
+                "`ray.tune.get_context()` when running in a function "
+                "passed to Ray Tune. This will be an error in the future. "
+                f"{V2_MIGRATION_GUIDE_MESSAGE}"
+            )
+        return get_tune_context()
+
+    global _default_context
+
+    with _context_lock:
+        if _default_context is None:
+            _default_context = TrainContext()
+        return _default_context

python/ray/train/data_parallel_trainer.py

@@ -0,0 +1,590 @@
+import logging
+import uuid
+from typing import Any, Callable, Dict, List, Optional, Type, Union
+
+import ray
+from ray._private.ray_constants import env_integer
+from ray._private.thirdparty.tabulate.tabulate import tabulate
+from ray.air.config import RunConfig, ScalingConfig
+from ray.train import BackendConfig, Checkpoint
+from ray.train._internal import session
+from ray.train._internal.backend_executor import BackendExecutor, TrialInfo
+from ray.train._internal.data_config import DataConfig
+from ray.train._internal.session import _TrainingResult, get_session
+from ray.train._internal.utils import construct_train_func, count_required_parameters
+from ray.train.base_trainer import _TRAINER_RESTORE_DEPRECATION_WARNING
+from ray.train.constants import RAY_TRAIN_ENABLE_STATE_TRACKING
+from ray.train.trainer import BaseTrainer, GenDataset, TrainingIterator
+from ray.util.annotations import Deprecated, DeveloperAPI
+from ray.widgets import Template
+from ray.widgets.util import repr_with_fallback
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class DataParallelTrainer(BaseTrainer):
+    """A Trainer for data parallel training.
+
+    You should subclass this Trainer if your Trainer follows SPMD (single program,
+    multiple data) programming paradigm - you want multiple processes to run the same
+    function, but on different data.
+
+    This Trainer runs the function ``train_loop_per_worker`` on multiple Ray
+    Actors.
+
+    The ``train_loop_per_worker`` function is expected to take in either 0 or 1
+    arguments:
+
+    .. testcode::
+
+        def train_loop_per_worker():
+            ...
+
+    .. testcode::
+
+        def train_loop_per_worker(config: Dict):
+            ...
+
+    If ``train_loop_per_worker`` accepts an argument, then
+    ``train_loop_config`` will be passed in as the argument. This is useful if you
+    want to tune the values in ``train_loop_config`` as hyperparameters.
+
+    If the ``datasets`` dict contains a training dataset (denoted by
+    the "train" key), then it will be split into multiple dataset
+    shards that can then be accessed by ``train.get_dataset_shard("train")`` inside
+    ``train_loop_per_worker``. All the other datasets will not be split and
+    ``train.get_dataset_shard(...)`` will return the entire Dataset.
+
+    Inside the ``train_loop_per_worker`` function, you can use any of the
+    :ref:`Ray Train loop methods <train-loop-api>`.
+
+    .. testcode::
+
+        from ray import train
+
+        def train_loop_per_worker():
+            # Report intermediate results for callbacks or logging and
+            # checkpoint data.
+            train.report(...)
+
+            # Returns dict of last saved checkpoint.
+            train.get_checkpoint()
+
+            # Returns the Dataset shard for the given key.
+            train.get_dataset_shard("my_dataset")
+
+            # Returns the total number of workers executing training.
+            train.get_context().get_world_size()
+
+            # Returns the rank of this worker.
+            train.get_context().get_world_rank()
+
+            # Returns the rank of the worker on the current node.
+            train.get_context().get_local_rank()
+
+    Any returns from the ``train_loop_per_worker`` will be discarded and not
+    used or persisted anywhere.
+
+    **How do I use DataParallelTrainer or any of its subclasses?**
+
+    Example:
+
+    .. testcode::
+        :skipif: True
+
+        import ray
+        from ray import train
+        from ray.train import ScalingConfig
+        from ray.train.data_parallel_trainer import DataParallelTrainer
+
+        def train_loop_for_worker():
+            dataset_shard_for_this_worker = train.get_dataset_shard("train")
+
+            # 3 items for 3 workers, each worker gets 1 item
+            batches = list(dataset_shard_for_this_worker.iter_batches(batch_size=1))
+            assert len(batches) == 1
+
+        train_dataset = ray.data.from_items([1, 2, 3])
+        assert train_dataset.count() == 3
+        trainer = DataParallelTrainer(
+            train_loop_for_worker,
+            scaling_config=ScalingConfig(num_workers=3),
+            datasets={"train": train_dataset},
+        )
+        result = trainer.fit()
+
+    **How do I develop on top of DataParallelTrainer?**
+
+    In many cases, using DataParallelTrainer directly is sufficient to execute
+    functions on multiple actors.
+
+    However, you may want to subclass ``DataParallelTrainer`` and create a custom
+    Trainer for the following 2 use cases:
+
+      - **Use Case 1:** You want to do data parallel training, but want to have
+        a predefined ``training_loop_per_worker``.
+
+      - **Use Case 2:** You want to implement a custom
+        :py:class:`~ray.train.backend.Backend` that automatically handles
+        additional setup or teardown logic on each actor, so that the users of this
+        new trainer do not have to implement this logic. For example, a
+        ``TensorflowTrainer`` can be built on top of ``DataParallelTrainer``
+        that automatically handles setting the proper environment variables for
+        distributed Tensorflow on each actor.
+
+    For 1, you can set a predefined training loop in __init__
+
+    .. testcode::
+
+        from ray.train.data_parallel_trainer import DataParallelTrainer
+
+        class MyDataParallelTrainer(DataParallelTrainer):
+            def __init__(self, *args, **kwargs):
+                predefined_train_loop_per_worker = lambda: 1
+                super().__init__(predefined_train_loop_per_worker, *args, **kwargs)
+
+
+    For 2, you can implement the ``ray.train.Backend`` and ``ray.train.BackendConfig``
+    interfaces.
+
+    .. testcode::
+
+        from dataclasses import dataclass
+        from ray.train.backend import Backend, BackendConfig
+
+        class MyBackend(Backend):
+            def on_start(self, worker_group, backend_config):
+                def set_env_var(env_var_value):
+                    import os
+                    os.environ["MY_ENV_VAR"] = env_var_value
+
+                worker_group.execute(set_env_var, backend_config.env_var)
+
+        @dataclass
+        class MyBackendConfig(BackendConfig):
+            env_var: str = "default_value"
+
+            def backend_cls(self):
+                return MyBackend
+
+        class MyTrainer(DataParallelTrainer):
+            def __init__(self, train_loop_per_worker, my_backend_config:
+                MyBackendConfig, **kwargs):
+
+                super().__init__(
+                    train_loop_per_worker,
+                    backend_config=my_backend_config, **kwargs)
+
+    Args:
+        train_loop_per_worker: The training function to execute.
+            This can either take in no arguments or a ``config`` dict.
+        train_loop_config: Configurations to pass into
+            ``train_loop_per_worker`` if it accepts an argument.
+        backend_config: Configuration for setting up a Backend (e.g. Torch,
+            Tensorflow, Horovod) on each worker to enable distributed
+            communication. If no Backend should be set up, then set this to None.
+        scaling_config: Configuration for how to scale data parallel training.
+        dataset_config: Configuration for dataset ingest. This is merged with the
+            default dataset config for the given trainer (`cls._dataset_config`).
+        run_config: Configuration for the execution of the training run.
+        datasets: Ray Datasets to use for training and evaluation.
+            This is a dict where the key is the name of the dataset, which
+            can be accessed from within the ``train_loop_per_worker`` by calling
+            ``train.get_dataset_shard(dataset_key)``.
+            By default, all datasets are sharded equally across workers.
+            This can be configured via ``dataset_config``.
+        metadata: Dict that should be made available via
+            `train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
+            for checkpoints saved from this Trainer. Must be JSON-serializable.
+        resume_from_checkpoint: A checkpoint to resume training from.
+    """
+
+    # Exposed here for testing purposes. Should never need
+    # to be overridden.
+    _backend_executor_cls: Type[BackendExecutor] = BackendExecutor
+    _training_iterator_cls: Type[TrainingIterator] = TrainingIterator
+
+    _scaling_config_allowed_keys = BaseTrainer._scaling_config_allowed_keys + [
+        "num_workers",
+        "resources_per_worker",
+        "use_gpu",
+        "placement_strategy",
+        "accelerator_type",
+    ]
+
+    # For backwards compatibility with the legacy dataset config API.
+    _dataset_config = None
+
+    _fields_for_tuner_param_space = BaseTrainer._fields_for_tuner_param_space + [
+        "train_loop_config"
+    ]
+
+    def __init__(
+        self,
+        train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]],
+        *,
+        train_loop_config: Optional[Dict] = None,
+        backend_config: Optional[BackendConfig] = None,
+        scaling_config: Optional[ScalingConfig] = None,
+        dataset_config: Optional[DataConfig] = None,
+        run_config: Optional[RunConfig] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        resume_from_checkpoint: Optional[Checkpoint] = None,
+    ):
+        self._train_loop_per_worker = train_loop_per_worker
+        self._train_loop_config = train_loop_config
+
+        if dataset_config is None:
+            dataset_config = DataConfig()
+
+        if not isinstance(dataset_config, DataConfig):
+            raise ValueError(
+                "`dataset_config` must be an instance of ray.train.DataConfig, "
+                f"was: {dataset_config}"
+            )
+        self._data_config = dataset_config
+
+        backend_config = (
+            backend_config if backend_config is not None else BackendConfig()
+        )
+        self._backend_config = backend_config
+
+        super(DataParallelTrainer, self).__init__(
+            scaling_config=scaling_config,
+            run_config=run_config,
+            datasets=datasets,
+            metadata=metadata,
+            resume_from_checkpoint=resume_from_checkpoint,
+        )
+
+        train_total_resources = self.scaling_config.total_resources
+        self._data_config.set_train_total_resources(
+            train_total_resources.get("CPU", 0),
+            train_total_resources.get("GPU", 0),
+        )
+
+        if env_integer(RAY_TRAIN_ENABLE_STATE_TRACKING, 0):
+            from ray.train._internal.state.state_actor import get_or_create_state_actor
+
+            get_or_create_state_actor()
+
+    @classmethod
+    @Deprecated(message=_TRAINER_RESTORE_DEPRECATION_WARNING)
+    def restore(
+        cls,
+        path: str,
+        train_loop_per_worker: Optional[
+            Union[Callable[[], None], Callable[[Dict], None]]
+        ] = None,
+        train_loop_config: Optional[Dict] = None,
+        **kwargs,
+    ):
+        """Restores a DataParallelTrainer from a previously interrupted/failed run.
+
+        Args:
+            path: The path to the experiment directory to restore from.
+            train_loop_per_worker: Optionally re-specified train loop function.
+                This should be used to re-specify a function that is not
+                restorable in a new Ray cluster (e.g., it holds onto outdated
+                object references). This should be the same training loop
+                that was passed to the original trainer constructor.
+            train_loop_config: Optionally re-specified train config.
+                This should similarly be used if the original `train_loop_config`
+                contained outdated object references, and it should not be modified
+                from what was originally passed in.
+            **kwargs: Additional arguments forwarded to
+                :meth:`BaseTrainer.restore() <ray.train.trainer.BaseTrainer.restore>`.
+
+        See :meth:`BaseTrainer.restore() <ray.train.trainer.BaseTrainer.restore>`
+        for descriptions of the other arguments.
+
+        Returns:
+            A restored instance of the ``DataParallelTrainer``.
+        """
+        return super(DataParallelTrainer, cls).restore(
+            path=path,
+            train_loop_per_worker=train_loop_per_worker,
+            train_loop_config=train_loop_config,
+            **kwargs,
+        )
+
+    def _validate_attributes(self):
+        super()._validate_attributes()
+
+        self._validate_train_loop_per_worker(
+            self._train_loop_per_worker, "train_loop_per_worker"
+        )
+
+    def _validate_train_loop_per_worker(
+        self, train_loop_per_worker: Callable, fn_name: str
+    ) -> None:
+        num_required_params = count_required_parameters(train_loop_per_worker)
+        if num_required_params > 1:
+            raise ValueError(
+                f"{fn_name} should take in 0 or 1 arguments, "
+                f"but it accepts {num_required_params} arguments instead."
+            )
+
+    @classmethod
+    def _validate_scaling_config(cls, scaling_config: ScalingConfig) -> ScalingConfig:
+        scaling_config = super(DataParallelTrainer, cls)._validate_scaling_config(
+            scaling_config
+        )
+
+        # This validation happens after the scaling config is updated from
+        # its specification in the Tuner `param_space`
+        if not scaling_config.use_gpu and "GPU" in ray.available_resources():
+            logger.info(
+                "GPUs are detected in your Ray cluster, but GPU "
+                "training is not enabled for this trainer. To enable "
+                "GPU training, make sure to set `use_gpu` to True "
+                "in your scaling config."
+            )
+
+        if scaling_config.num_workers is None:
+            raise ValueError(
+                "You must specify the 'num_workers' in `scaling_config` as either an "
+                f"argument of `{cls.__name__}` or through the `param_space` of a "
+                "`Tuner` (if performing hyperparameter tuning)."
+            )
+
+        if scaling_config.num_workers <= 0:
+            raise ValueError(
+                "'num_workers' in `scaling_config` must be a positive "
+                f"integer. Received {scaling_config.num_workers}"
+            )
+
+        return scaling_config
+
+    def _run_training(self, training_iterator: TrainingIterator) -> None:
+        """This method loops over the `TrainingIterator`:
+        The actual iteration (for ... in ...) waits for the training function
+        on each worker to report a result and supplies it as a list of results.
+        Afterwards (in the body of the loop), it will report the result
+        to the Tune session.
+        The iterator ends after the training function on each worker has finished.
+        """
+        for training_results in training_iterator:
+            # TODO(ml-team): add ability to report results from multiple workers.
+            self._propagate_results(training_results)
+
+    def _propagate_results(self, training_results: List[_TrainingResult]):
+        first_worker_result = training_results[0]
+        assert all(isinstance(result, _TrainingResult) for result in training_results)
+
+        tune_session = get_session()
+
+        # Check if any workers reported a checkpoint.
+        # If so, report a checkpoint pointing to the persisted location
+        # to Tune for book-keeping.
+        # NOTE: This removes the restriction for any individual worker
+        # (ex: global rank 0 worker) from needing to report a checkpoint.
+        # All workers reported a checkpoint to the same fs path, so there's
+        # no need to report multiple checkpoints to Tune.
+        worker_checkpoints = [
+            result.checkpoint
+            for result in training_results
+            if result.checkpoint is not None
+        ]
+        at_least_one_reported_checkpoint = len(worker_checkpoints) > 0
+
+        if at_least_one_reported_checkpoint:
+            # Update the coordinator's checkpoint index to the latest.
+            # This is what keeps the checkpoint index in line with the workers.
+            tune_session.storage._update_checkpoint_index(first_worker_result.metrics)
+
+        # Make sure that all workers uploaded to the same location.
+        assert all(
+            checkpoint.path == tune_session.storage.checkpoint_fs_path
+            for checkpoint in worker_checkpoints
+        )
+
+        checkpoint = (
+            Checkpoint(
+                filesystem=tune_session.storage.storage_filesystem,
+                path=tune_session.storage.checkpoint_fs_path,
+            )
+            if at_least_one_reported_checkpoint
+            else None
+        )
+
+        tracked_training_result = _TrainingResult(
+            checkpoint=checkpoint,
+            metrics=first_worker_result.metrics,
+        )
+
+        logger.debug(
+            "Report (metrics, checkpoint) to the Tune session:\n"
+            f"  metrics={tracked_training_result.metrics}\n"
+            f"  checkpoint={tracked_training_result.checkpoint}"
+        )
+
+        # Report the metrics and checkpoint to Tune.
+        tune_session._report_training_result(tracked_training_result)
+
+    def training_loop(self) -> None:
+        scaling_config = self._validate_scaling_config(self.scaling_config)
+
+        train_loop_per_worker = construct_train_func(
+            self._train_loop_per_worker,
+            self._train_loop_config,
+            train_func_context=self._backend_config.train_func_context,
+            fn_arg_name="train_loop_per_worker",
+            discard_returns=True,
+        )
+
+        trial_info = TrialInfo(
+            name=session.get_trial_name(),
+            id=session.get_trial_id(),
+            resources=session.get_trial_resources(),
+            logdir=session.get_trial_dir(),
+            driver_ip=ray.util.get_node_ip_address(),
+            driver_node_id=ray.get_runtime_context().get_node_id(),
+            experiment_name=session.get_experiment_name(),
+            run_id=uuid.uuid4().hex,
+        )
+
+        backend_executor = self._backend_executor_cls(
+            backend_config=self._backend_config,
+            trial_info=trial_info,
+            num_workers=scaling_config.num_workers,
+            resources_per_worker=scaling_config._resources_per_worker_not_none,
+            max_retries=0,
+        )
+
+        # Start the remote actors.
+        backend_executor.start()
+
+        training_iterator = self._training_iterator_cls(
+            backend_executor=backend_executor,
+            backend_config=self._backend_config,
+            train_func=train_loop_per_worker,
+            datasets=self.datasets,
+            metadata=self.metadata,
+            data_config=self._data_config,
+            checkpoint=self.starting_checkpoint,
+        )
+
+        self._run_training(training_iterator)
+
+        # Shutdown workers.
+        backend_executor.shutdown()
+
+    def get_dataset_config(self) -> DataConfig:
+        """Returns a copy of this Trainer's final dataset configs.
+
+        Returns:
+            The merged default + user-supplied dataset config.
+        """
+
+        return self._data_config
+
+    @repr_with_fallback(["ipywidgets", "8"])
+    def _repr_mimebundle_(self, **kwargs: Any):
+        """Returns a mimebundle with an ipywidget repr and a simple text repr.
+
+        Depending on the frontend where the data is being displayed,
+        different mimetypes will be used from this bundle.
+        See https://ipython.readthedocs.io/en/stable/config/integrating.html
+        for information about this method, and
+        https://ipywidgets.readthedocs.io/en/latest/embedding.html
+        for more information about the jupyter widget mimetype.
+
+        Args:
+            **kwargs: Standard Jupyter mimebundle kwargs (e.g. ``include``,
+                ``exclude``); unused by this implementation.
+
+        Returns:
+            A mimebundle containing an ipywidget repr and a simple text repr.
+        """
+        from ipywidgets import HTML, Layout, Tab, VBox
+
+        title = HTML(f"<h2>{self.__class__.__name__}</h2>")
+
+        children = []
+        titles = []
+
+        if self.datasets:
+            children.append(self._datasets_repr_())
+            titles.append("Datasets")
+
+            children.append(HTML(self._data_config_repr_html_()))
+            titles.append("Data Config")
+
+        if self._train_loop_config:
+            children.append(HTML(self._train_loop_config_repr_html_()))
+            titles.append("Train Loop Config")
+
+        if self.scaling_config:
+            children.append(HTML(self.scaling_config._repr_html_()))
+            titles.append("Scaling Config")
+
+        if self.run_config:
+            children.append(HTML(self.run_config._repr_html_()))
+            titles.append("Run Config")
+
+        if self._backend_config:
+            children.append(HTML(self._backend_config._repr_html_()))
+            titles.append("Backend Config")
+
+        tab = Tab(children, titles=titles)
+        widget = VBox([title, tab], layout=Layout(width="100%"))
+        bundle = widget._repr_mimebundle_(**kwargs)
+        bundle.update(
+            {
+                "text/plain": repr(self),
+            }
+        )
+        return bundle
+
+    def _train_loop_config_repr_html_(self) -> str:
+        if self._train_loop_config:
+            table_data = {}
+            for k, v in self._train_loop_config.items():
+                if isinstance(v, str) or str(v).isnumeric():
+                    table_data[k] = v
+                elif hasattr(v, "_repr_html_"):
+                    table_data[k] = v._repr_html_()
+                else:
+                    table_data[k] = str(v)
+
+            return Template("title_data.html.j2").render(
+                title="Train Loop Config",
+                data=Template("scrollableTable.html.j2").render(
+                    table=tabulate(
+                        table_data.items(),
+                        headers=["Setting", "Value"],
+                        showindex=False,
+                        tablefmt="unsafehtml",
+                    ),
+                    max_height="none",
+                ),
+            )
+        else:
+            return ""
+
+    def _data_config_repr_html_(self) -> str:
+        # TODO make this rendering nicer.
+        content = [str(self._data_config)]
+        return Template("rendered_html_common.html.j2").render(content=content)
+
+    def _datasets_repr_(self) -> str:
+        from ipywidgets import HTML, Layout, VBox
+
+        content = []
+        if self.datasets:
+            for name, config in self.datasets.items():
+                tab = config._tab_repr_()
+                if tab:
+                    content.append(
+                        HTML(
+                            Template("title_data.html.j2").render(
+                                title=f"Dataset - <code>{name}</code>", data=None
+                            )
+                        )
+                    )
+                    content.append(config._tab_repr_())
+
+        return VBox(content, layout=Layout(width="100%"))

python/ray/train/error.py

@@ -0,0 +1,6 @@
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="beta")
+class SessionMisuseError(Exception):
+    """Indicates a method or function was used outside of a session."""

python/ray/train/examples/accelerate/accelerate_torch_trainer.py

@@ -0,0 +1,164 @@
+# __accelerate_torch_basic_example_start__
+"""
+Minimal Ray Train and Accelerate example adapted from
+https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
+
+Fine-tune a BERT model with Hugging Face Accelerate and Ray Train and Ray Data
+"""
+
+from tempfile import TemporaryDirectory
+
+import evaluate
+import torch
+from accelerate import Accelerator
+from datasets import load_dataset
+from torch.optim import AdamW
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+import ray
+import ray.train
+from ray.train import Checkpoint, DataConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    """Your training function that launches on each worker."""
+
+    # Unpack training configs
+    lr = config["lr"]
+    seed = config["seed"]
+    num_epochs = config["num_epochs"]
+    train_batch_size = config["train_batch_size"]
+    eval_batch_size = config["eval_batch_size"]
+    train_ds_size = config["train_dataset_size"]
+
+    set_seed(seed)
+
+    # Initialize accelerator
+    accelerator = Accelerator()
+
+    # Load datasets and metrics
+    metric = evaluate.load("glue", "mrpc")
+
+    # Prepare Ray Data loaders
+    # ====================================================
+    train_ds = ray.train.get_dataset_shard("train")
+    eval_ds = ray.train.get_dataset_shard("validation")
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def collate_fn(batch):
+        outputs = tokenizer(
+            list(batch["sentence1"]),
+            list(batch["sentence2"]),
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+        )
+        outputs["labels"] = torch.LongTensor(batch["label"])
+        outputs = {k: v.to(accelerator.device) for k, v in outputs.items()}
+        return outputs
+
+    train_dataloader = train_ds.iter_torch_batches(
+        batch_size=train_batch_size, collate_fn=collate_fn
+    )
+    eval_dataloader = eval_ds.iter_torch_batches(
+        batch_size=eval_batch_size, collate_fn=collate_fn
+    )
+    # ====================================================
+
+    # Instantiate the model, optimizer, lr_scheduler
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", return_dict=True
+    )
+
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    steps_per_epoch = train_ds_size // (accelerator.num_processes * train_batch_size)
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(steps_per_epoch * num_epochs),
+    )
+
+    # Prepare everything with accelerator
+    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        # Evaluation
+        model.eval()
+        for batch in eval_dataloader:
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+
+            predictions, references = accelerator.gather_for_metrics(
+                (predictions, batch["labels"])
+            )
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # Report checkpoint and metrics to Ray Train
+        # ==========================================
+        with TemporaryDirectory() as tmpdir:
+            if accelerator.is_main_process:
+                unwrapped_model = accelerator.unwrap_model(model)
+                accelerator.save(unwrapped_model, f"{tmpdir}/ckpt_{epoch}.bin")
+                checkpoint = Checkpoint.from_directory(tmpdir)
+            else:
+                checkpoint = None
+            ray.train.report(metrics=eval_metric, checkpoint=checkpoint)
+
+
+if __name__ == "__main__":
+    config = {
+        "lr": 2e-5,
+        "num_epochs": 3,
+        "seed": 42,
+        "train_batch_size": 16,
+        "eval_batch_size": 32,
+    }
+
+    # Prepare Ray Datasets
+    hf_datasets = load_dataset("nyu-mll/glue", "mrpc")
+    ray_datasets = {
+        "train": ray.data.from_huggingface(hf_datasets["train"]),
+        "validation": ray.data.from_huggingface(hf_datasets["validation"]),
+    }
+    config["train_dataset_size"] = ray_datasets["train"].count()
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=config,
+        datasets=ray_datasets,
+        dataset_config=DataConfig(datasets_to_split=["train", "validation"]),
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        # If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        # run_config=ray.train.RunConfig(storage_path="s3://..."),
+    )
+
+    result = trainer.fit()
+
+# __accelerate_torch_basic_example_end__

python/ray/train/examples/accelerate/accelerate_torch_trainer_no_raydata.py

@@ -0,0 +1,169 @@
+# __accelerate_torch_basic_example_no_raydata_start__
+"""
+Minimal Ray Train + Accelerate example adapted from
+https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
+
+Fine-tune a BERT model with Hugging Face Accelerate and Ray Train
+"""
+
+from tempfile import TemporaryDirectory
+
+import evaluate
+import torch
+from accelerate import Accelerator
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+import ray.train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    """Your training function that will be launched on each worker."""
+
+    # Unpack training configs
+    lr = config["lr"]
+    seed = config["seed"]
+    num_epochs = config["num_epochs"]
+    train_batch_size = config["train_batch_size"]
+    eval_batch_size = config["eval_batch_size"]
+
+    set_seed(seed)
+
+    # Initialize accelerator
+    accelerator = Accelerator()
+
+    # Load datasets and metrics
+    metric = evaluate.load("glue", "mrpc")
+
+    # Prepare PyTorch DataLoaders
+    # ====================================================
+    hf_datasets = load_dataset("nyu-mll/glue", "mrpc")
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def collate_fn(batch):
+        outputs = tokenizer(
+            [sample["sentence1"] for sample in batch],
+            [sample["sentence2"] for sample in batch],
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+        )
+        outputs["labels"] = torch.LongTensor([sample["label"] for sample in batch])
+        outputs = {k: v.to(accelerator.device) for k, v in outputs.items()}
+        return outputs
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        hf_datasets["train"],
+        shuffle=True,
+        collate_fn=collate_fn,
+        batch_size=train_batch_size,
+        drop_last=True,
+    )
+    eval_dataloader = DataLoader(
+        hf_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=eval_batch_size,
+        drop_last=True,
+    )
+    # ====================================================
+
+    # Instantiate the model, optimizer, lr_scheduler
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", return_dict=True
+    )
+
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    steps_per_epoch = len(train_dataloader)
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(steps_per_epoch * num_epochs),
+    )
+
+    # Prepare everything with accelerator
+    (
+        model,
+        optimizer,
+        train_dataloader,
+        eval_dataloader,
+        lr_scheduler,
+    ) = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        # Evaluation
+        model.eval()
+        for batch in eval_dataloader:
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+
+            predictions, references = accelerator.gather_for_metrics(
+                (predictions, batch["labels"])
+            )
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # Report Checkpoint and metrics to Ray Train
+        # ==========================================
+        with TemporaryDirectory() as tmpdir:
+            if accelerator.is_main_process:
+                unwrapped_model = accelerator.unwrap_model(model)
+                accelerator.save(unwrapped_model, f"{tmpdir}/ckpt_{epoch}.bin")
+                checkpoint = Checkpoint.from_directory(tmpdir)
+            else:
+                checkpoint = None
+            ray.train.report(metrics=eval_metric, checkpoint=checkpoint)
+
+
+if __name__ == "__main__":
+    config = {
+        "lr": 2e-5,
+        "num_epochs": 3,
+        "seed": 42,
+        "train_batch_size": 16,
+        "eval_batch_size": 32,
+    }
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        # If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        # run_config=ray.train.RunConfig(storage_path="s3://..."),
+    )
+
+    result = trainer.fit()
+
+# __accelerate_torch_basic_example_no_raydata_end__

python/ray/train/examples/deepspeed/deepspeed_torch_trainer.py

@@ -0,0 +1,185 @@
+# __deepspeed_torch_basic_example_start__
+"""
+Minimal Ray Train + DeepSpeed example adapted from
+https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
+
+Fine-tune a BERT model with DeepSpeed ZeRO-3 and Ray Train and Ray Data
+"""
+
+from tempfile import TemporaryDirectory
+
+import deepspeed
+import torch
+from datasets import load_dataset
+from deepspeed.accelerator import get_accelerator
+from torchmetrics.classification import BinaryAccuracy, BinaryF1Score
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, set_seed
+
+import ray
+import ray.train
+from ray.train import Checkpoint, DataConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    """Your training function that will be launched on each worker."""
+
+    # Unpack training configs
+    set_seed(config["seed"])
+    num_epochs = config["num_epochs"]
+    train_batch_size = config["train_batch_size"]
+    eval_batch_size = config["eval_batch_size"]
+
+    # Instantiate the Model
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", return_dict=True
+    )
+
+    # Prepare Ray Data Loaders
+    # ====================================================
+    train_ds = ray.train.get_dataset_shard("train")
+    eval_ds = ray.train.get_dataset_shard("validation")
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def collate_fn(batch):
+        outputs = tokenizer(
+            list(batch["sentence1"]),
+            list(batch["sentence2"]),
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+        )
+        outputs["labels"] = torch.LongTensor(batch["label"])
+        return outputs
+
+    train_dataloader = train_ds.iter_torch_batches(
+        batch_size=train_batch_size, collate_fn=collate_fn
+    )
+    eval_dataloader = eval_ds.iter_torch_batches(
+        batch_size=eval_batch_size, collate_fn=collate_fn
+    )
+    # ====================================================
+
+    # Initialize DeepSpeed Engine
+    model, optimizer, _, lr_scheduler = deepspeed.initialize(
+        model=model,
+        model_parameters=model.parameters(),
+        config=deepspeed_config,
+    )
+    device = get_accelerator().device_name(model.local_rank)
+
+    # Initialize Evaluation Metrics
+    f1 = BinaryF1Score().to(device)
+    accuracy = BinaryAccuracy().to(device)
+
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        for batch in train_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            outputs = model(**batch)
+            loss = outputs.loss
+            model.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        # Evaluation
+        model.eval()
+        for batch in eval_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+
+            f1.update(predictions, batch["labels"])
+            accuracy.update(predictions, batch["labels"])
+
+        # torchmetrics will aggregate the metrics across all workers
+        eval_metric = {
+            "f1": f1.compute().item(),
+            "accuracy": accuracy.compute().item(),
+        }
+        f1.reset()
+        accuracy.reset()
+
+        if model.global_rank == 0:
+            print(f"epoch {epoch}:", eval_metric)
+
+        # Report checkpoint and metrics to Ray Train
+        # ==============================================================
+        with TemporaryDirectory() as tmpdir:
+            # Each worker saves its own checkpoint shard
+            model.save_checkpoint(tmpdir)
+
+            # Ensure all workers finished saving their checkpoint shard
+            torch.distributed.barrier()
+
+            # Report checkpoint shards from each worker in parallel
+            ray.train.report(
+                metrics=eval_metric, checkpoint=Checkpoint.from_directory(tmpdir)
+            )
+        # ==============================================================
+
+
+if __name__ == "__main__":
+    deepspeed_config = {
+        "optimizer": {
+            "type": "AdamW",
+            "params": {
+                "lr": 2e-5,
+            },
+        },
+        "scheduler": {"type": "WarmupLR", "params": {"warmup_num_steps": 100}},
+        "fp16": {"enabled": True},
+        "bf16": {"enabled": False},  # Turn this on if using AMPERE GPUs.
+        "zero_optimization": {
+            "stage": 3,
+            "offload_optimizer": {
+                "device": "none",
+            },
+            "offload_param": {
+                "device": "none",
+            },
+        },
+        "gradient_accumulation_steps": 1,
+        "gradient_clipping": True,
+        "steps_per_print": 10,
+        "train_micro_batch_size_per_gpu": 16,
+        "wall_clock_breakdown": False,
+    }
+
+    training_config = {
+        "seed": 42,
+        "num_epochs": 3,
+        "train_batch_size": 16,
+        "eval_batch_size": 32,
+        "deepspeed_config": deepspeed_config,
+    }
+
+    # Prepare Ray Datasets
+    hf_datasets = load_dataset("nyu-mll/glue", "mrpc")
+    ray_datasets = {
+        "train": ray.data.from_huggingface(hf_datasets["train"]),
+        "validation": ray.data.from_huggingface(hf_datasets["validation"]),
+    }
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=training_config,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        datasets=ray_datasets,
+        dataset_config=DataConfig(datasets_to_split=["train", "validation"]),
+        # If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        # run_config=ray.train.RunConfig(storage_path="s3://..."),
+    )
+
+    result = trainer.fit()
+
+    # Retrieve the best checkponints from results
+    _ = result.best_checkpoints
+
+# __deepspeed_torch_basic_example_end__

python/ray/train/examples/deepspeed/deepspeed_torch_trainer_no_raydata.py

@@ -0,0 +1,178 @@
+# __deepspeed_torch_basic_example_no_raydata_start__
+"""
+Minimal Ray Train + DeepSpeed example adapted from
+https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
+
+Fine-tune a BERT model with DeepSpeed ZeRO-3 and Ray Train
+"""
+
+from tempfile import TemporaryDirectory
+
+import deepspeed
+import torch
+from datasets import load_dataset
+from deepspeed.accelerator import get_accelerator
+from torch.utils.data import DataLoader
+from torchmetrics.classification import BinaryAccuracy, BinaryF1Score
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, set_seed
+
+import ray
+import ray.train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    """Your training function that will be launched on each worker."""
+
+    # Unpack training configs
+    set_seed(config["seed"])
+    num_epochs = config["num_epochs"]
+    eval_batch_size = config["eval_batch_size"]
+
+    # Instantiate the Model
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", return_dict=True
+    )
+
+    # Prepare PyTorch Data Loaders
+    # ====================================================
+    hf_datasets = load_dataset("nyu-mll/glue", "mrpc")
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def collate_fn(batch):
+        outputs = tokenizer(
+            [sample["sentence1"] for sample in batch],
+            [sample["sentence2"] for sample in batch],
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+        )
+        outputs["labels"] = torch.LongTensor([sample["label"] for sample in batch])
+        return outputs
+
+    # Instantiate dataloaders.
+    # The train_dataloader already created by `deepspeed.initialize`
+    eval_dataloader = DataLoader(
+        hf_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=eval_batch_size,
+        drop_last=True,
+    )
+    # ====================================================
+
+    # Initialize DeepSpeed Engine
+    model, optimizer, train_dataloader, lr_scheduler = deepspeed.initialize(
+        model=model,
+        model_parameters=model.parameters(),
+        training_data=hf_datasets["train"],
+        collate_fn=collate_fn,
+        config=deepspeed_config,
+    )
+    device = get_accelerator().device_name(model.local_rank)
+
+    # Initialize Evaluation Metrics
+    f1 = BinaryF1Score().to(device)
+    accuracy = BinaryAccuracy().to(device)
+
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        for batch in train_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            outputs = model(**batch)
+            loss = outputs.loss
+            model.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        # Evaluation
+        model.eval()
+        for batch in eval_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+
+            f1.update(predictions, batch["labels"])
+            accuracy.update(predictions, batch["labels"])
+
+        # torchmetrics will aggregate the metrics across all workers
+        eval_metric = {
+            "f1": f1.compute().item(),
+            "accuracy": accuracy.compute().item(),
+        }
+        f1.reset()
+        accuracy.reset()
+
+        if model.global_rank == 0:
+            print(f"epoch {epoch}:", eval_metric)
+
+        # Report checkpoint and metrics to Ray Train
+        # ==============================================================
+        with TemporaryDirectory() as tmpdir:
+            # Each worker saves its own checkpoint shard
+            model.save_checkpoint(tmpdir)
+
+            # Ensure all workers finished saving their checkpoint shard
+            torch.distributed.barrier()
+
+            # Report checkpoint shards from each worker in parallel
+            ray.train.report(
+                metrics=eval_metric, checkpoint=Checkpoint.from_directory(tmpdir)
+            )
+        # ==============================================================
+
+
+if __name__ == "__main__":
+    deepspeed_config = {
+        "optimizer": {
+            "type": "AdamW",
+            "params": {
+                "lr": 2e-5,
+            },
+        },
+        "scheduler": {"type": "WarmupLR", "params": {"warmup_num_steps": 100}},
+        "fp16": {"enabled": True},
+        "bf16": {"enabled": False},  # Turn this on if using AMPERE GPUs.
+        "zero_optimization": {
+            "stage": 3,
+            "offload_optimizer": {
+                "device": "none",
+            },
+            "offload_param": {
+                "device": "none",
+            },
+        },
+        "gradient_accumulation_steps": 1,
+        "gradient_clipping": True,
+        "steps_per_print": 10,
+        "train_micro_batch_size_per_gpu": 16,
+        "wall_clock_breakdown": False,
+    }
+
+    training_config = {
+        "seed": 42,
+        "num_epochs": 3,
+        "eval_batch_size": 32,
+        "deepspeed_config": deepspeed_config,
+    }
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=training_config,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        # If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        # run_config=ray.train.RunConfig(storage_path="s3://..."),
+    )
+
+    result = trainer.fit()
+
+    # Retrieve the best checkponints from results
+    _ = result.best_checkpoints
+
+# __deepspeed_torch_basic_example_no_raydata_end__

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_comet.py

@@ -0,0 +1,41 @@
+# isort: skip_file
+from lightning_exp_tracking_model_dl import DummyModel, dataloader
+
+# __lightning_experiment_tracking_comet_start__
+import os
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import CometLogger
+
+
+def train_func(config):
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = CometLogger(api_key=os.environ["COMET_API_KEY"])
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "COMET_API_KEY" in os.environ
+), 'Please do COMET_API_KEY="abcde" when running this script.'
+# This makes sure that all workers have this env var set.
+ray.init(runtime_env={"env_vars": {"COMET_API_KEY": os.environ["COMET_API_KEY"]}})
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=scaling_config,
+)
+
+trainer.fit()

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_mlflow.py

@@ -0,0 +1,63 @@
+# ruff: noqa
+# isort: skip_file
+import os
+import tempfile
+
+tempdir = tempfile.TemporaryDirectory()
+os.environ["SHARED_STORAGE_PATH"] = tempdir.name
+
+from ray.train.examples.experiment_tracking.lightning_exp_tracking_model_dl import (
+    DummyModel,
+    dataloader,
+)
+
+
+# __lightning_experiment_tracking_mlflow_start__
+import os
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import MLFlowLogger
+
+
+def train_func(config):
+
+    save_dir = config["save_dir"]
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = MLFlowLogger(
+            experiment_name="demo-project",
+            tracking_uri=f"file:{save_dir}",
+        )
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "SHARED_STORAGE_PATH" in os.environ
+), "Please do SHARED_STORAGE_PATH=/a/b/c when running this script."
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={
+        "save_dir": os.path.join(os.environ["SHARED_STORAGE_PATH"], "mlruns")
+    },
+    scaling_config=scaling_config,
+)
+
+trainer.fit()
+# __lightning_experiment_tracking_mlflow_end__
+
+tempdir.cleanup()

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_model_dl.py

@@ -0,0 +1,46 @@
+# ruff: noqa
+# fmt: off
+# # isort: skip_file
+
+
+# __model_dl_start__
+import lightning.pytorch as pl
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+
+# Create dummy data
+X = torch.randn(128, 3)  # 128 samples, 3 features
+y = torch.randint(0, 2, (128,))  # 128 binary labels
+
+# Create a TensorDataset to wrap the data
+dataset = TensorDataset(X, y)
+
+# Create a DataLoader to iterate over the dataset
+batch_size = 8
+dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+
+# Define a dummy model
+class DummyModel(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.layer = torch.nn.Linear(3, 1)
+
+    def forward(self, x):
+        return self.layer(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        loss = F.binary_cross_entropy_with_logits(y_hat.flatten(), y.float())
+
+        # The metrics below will be reported to Loggers
+        self.log("train_loss", loss)
+        self.log_dict({
+            "metric_1": 1 / (batch_idx + 1), "metric_2": batch_idx * 100
+        })
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=1e-3)

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_tensorboard.py

@@ -0,0 +1,60 @@
+# ruff: noqa
+# isort: skip_file
+
+import os
+import tempfile
+
+tempdir = tempfile.TemporaryDirectory()
+os.environ["SHARED_STORAGE_PATH"] = tempdir.name
+
+from ray.train.examples.experiment_tracking.lightning_exp_tracking_model_dl import (
+    DummyModel,
+    dataloader,
+)
+
+# __lightning_experiment_tracking_tensorboard_start__
+import os
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import TensorBoardLogger
+
+
+def train_func(config):
+
+    save_dir = config["save_dir"]
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = TensorBoardLogger(name="demo-run", save_dir=f"file:{save_dir}")
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "SHARED_STORAGE_PATH" in os.environ
+), "Please do SHARED_STORAGE_PATH=/a/b/c when running this script."
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={
+        "save_dir": os.path.join(os.environ["SHARED_STORAGE_PATH"], "tensorboard")
+    },
+    scaling_config=scaling_config,
+)
+
+trainer.fit()
+# __lightning_experiment_tracking_tensorboard_end__
+
+tempdir.cleanup()

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_wandb.py

@@ -0,0 +1,50 @@
+# ruff: noqa
+# fmt: off
+# # isort: skip_file
+
+from lightning_exp_tracking_model_dl import DummyModel, dataloader
+
+# __lightning_experiment_tracking_wandb_start__
+import os
+import wandb
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import WandbLogger
+
+
+def train_func(config):
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = WandbLogger(name="demo-run", project="demo-project")
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.finish()
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "WANDB_API_KEY" in os.environ
+), 'Please set WANDB_API_KEY="abcde" when running this script.'
+
+# This ensures that all workers have this env var set.
+ray.init(
+    runtime_env={"env_vars": {"WANDB_API_KEY": os.environ["WANDB_API_KEY"]}}
+)
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=scaling_config,
+)
+
+trainer.fit()

python/ray/train/examples/experiment_tracking/torch_exp_tracking_mlflow.py

@@ -0,0 +1,85 @@
+# ruff: noqa
+# isort: skip_file
+from filelock import FileLock
+import os
+import tempfile
+
+tempdir = tempfile.TemporaryDirectory()
+os.environ["SHARED_STORAGE_PATH"] = tempdir.name
+
+# __start__
+# Run the following script with the SHARED_STORAGE_PATH env var set.
+# The MLflow offline logs are saved to SHARED_STORAGE_PATH/mlruns.
+
+import mlflow
+import os
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+import torch
+from torchvision import datasets, transforms
+from torchvision.models import resnet18
+from torch.utils.data import DataLoader
+
+assert os.environ.get(
+    "SHARED_STORAGE_PATH", None
+), "Please set SHARED_STORAGE_PATH env var."
+
+
+# Assumes you are passing a `save_dir` in `config`
+def train_func(config):
+    save_dir = config["save_dir"]
+    if ray.train.get_context().get_world_rank() == 0:
+        mlflow.set_tracking_uri(f"file:{save_dir}")
+        mlflow.set_experiment("my_experiment")
+        mlflow.start_run()
+
+    # Model, Loss, Optimizer
+    model = resnet18(num_classes=10)
+    model.conv1 = torch.nn.Conv2d(
+        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
+    )
+    model = ray.train.torch.prepare_model(model)
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.module.parameters(), lr=0.001)
+
+    # Data
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.28604,), (0.32025,))]
+    )
+    with FileLock("./data.lock"):
+        train_data = datasets.FashionMNIST(
+            root="./data", train=True, download=True, transform=transform
+        )
+    train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
+    train_loader = ray.train.torch.prepare_data_loader(train_loader)
+
+    # Training
+    for epoch in range(1):
+        if ray.train.get_context().get_world_size() > 1:
+            train_loader.sampler.set_epoch(epoch)
+
+        for images, labels in train_loader:
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if ray.train.get_context().get_world_rank() == 0:
+                mlflow.log_metrics({"loss": loss.item(), "epoch": epoch})
+
+    if ray.train.get_context().get_world_rank() == 0:
+        mlflow.end_run()
+
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={
+        "save_dir": os.path.join(os.environ["SHARED_STORAGE_PATH"], "mlruns")
+    },
+    scaling_config=ScalingConfig(num_workers=2),
+)
+trainer.fit()
+# __end__
+
+tempdir.cleanup()

python/ray/train/examples/experiment_tracking/torch_exp_tracking_wandb.py

@@ -0,0 +1,75 @@
+# ruff: noqa
+# fmt: off
+# isort: off
+
+# __start__
+from filelock import FileLock
+import os
+
+import torch
+import wandb
+
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from torchvision.models import resnet18
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+# Run the following script with the WANDB_API_KEY env var set.
+assert os.environ.get("WANDB_API_KEY", None), "Please set WANDB_API_KEY env var."
+
+# This makes sure that all workers have this env var set.
+ray.init(
+    runtime_env={"env_vars": {"WANDB_API_KEY": os.environ["WANDB_API_KEY"]}}
+)
+
+
+def train_func(config):
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.init()
+
+    # Model, Loss, Optimizer
+    model = resnet18(num_classes=10)
+    model.conv1 = torch.nn.Conv2d(
+        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
+    )
+    model = ray.train.torch.prepare_model(model)
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.module.parameters(), lr=0.001)
+
+    # Data
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.28604,), (0.32025,))]
+    )
+    with FileLock("./data.lock"):
+        train_data = datasets.FashionMNIST(
+            root="./data", train=True, download=True, transform=transform
+        )
+    train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
+    train_loader = ray.train.torch.prepare_data_loader(train_loader)
+
+    # Training
+    for epoch in range(1):
+        if ray.train.get_context().get_world_size() > 1:
+            train_loader.sampler.set_epoch(epoch)
+
+        for images, labels in train_loader:
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if ray.train.get_context().get_world_rank() == 0:
+                wandb.log({"loss": loss, "epoch": epoch})
+
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.finish()
+
+
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=ScalingConfig(num_workers=2),
+)
+trainer.fit()

python/ray/train/examples/horovod/horovod-cluster.yaml

@@ -0,0 +1,55 @@
+# An unique identifier for the head node and workers of this cluster.
+cluster_name: horovod-cluster
+
+# The maximum number of workers nodes to launch in addition to the head
+# node. This takes precedence over min_workers. min_workers default to 0.
+min_workers: 3
+max_workers: 3
+
+# Cloud-provider specific configuration.
+provider:
+    type: aws
+    region: us-west-2
+
+# How Ray will authenticate with newly launched nodes.
+auth:
+    ssh_user: ubuntu
+
+available_node_types:
+    ray.head.default:
+        min_workers: 0
+        max_workers: 0
+        resources: {}
+        node_config:
+            InstanceType: g3.8xlarge
+            ImageId: latest_dlami
+            InstanceMarketOptions:
+                MarketType: spot
+            BlockDeviceMappings:
+                - DeviceName: /dev/sda1
+                  Ebs:
+                      VolumeSize: 300
+
+
+    ray.worker.default:
+        min_workers: 3
+        max_workers: 3
+        resources: {}
+        node_config:
+            InstanceType: g3.8xlarge
+            ImageId: latest_dlami
+            InstanceMarketOptions:
+                MarketType: spot
+            BlockDeviceMappings:
+                - DeviceName: /dev/sda1
+                  Ebs:
+                      VolumeSize: 300
+
+
+setup_commands:
+    # This replaces the standard anaconda Ray installation
+    - pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl
+    - pip install ray[tune]
+
+    # Install Horovod
+    - HOROVOD_WITH_GLOO=1 HOROVOD_GPU_OPERATIONS=NCCL HOROVOD_WITHOUT_MPI=1 HOROVOD_WITHOUT_TENSORFLOW=1 HOROVOD_WITHOUT_MXNET=1 HOROVOD_WITH_PYTORCH=1 pip install torch torchvision horovod

python/ray/train/examples/horovod/horovod_example.py

@@ -0,0 +1,286 @@
+import argparse
+import os
+
+import horovod.torch as hvd
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.utils.data.distributed
+from filelock import FileLock
+from torchvision import datasets, transforms
+
+import ray
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.horovod import HorovodTrainer
+
+
+def metric_average(val, name):
+    tensor = torch.tensor(val)
+    avg_tensor = hvd.allreduce(tensor, name=name)
+    return avg_tensor.item()
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x)
+
+
+def setup(config):
+    data_dir = config.get("data_dir", None)
+    seed = config.get("seed", 42)
+    batch_size = config.get("batch_size", 64)
+    use_adasum = config.get("use_adasum", False)
+    lr = config.get("lr", 0.01)
+    momentum = config.get("momentum", 0.5)
+    use_cuda = config.get("use_cuda", False)
+
+    # Horovod: initialize library.
+    hvd.init()
+    torch.manual_seed(seed)
+
+    if use_cuda:
+        # Horovod: pin GPU to local rank.
+        torch.cuda.set_device(hvd.local_rank())
+        torch.cuda.manual_seed(seed)
+
+    # Horovod: limit # of CPU threads to be used per worker.
+    torch.set_num_threads(1)
+
+    kwargs = {"num_workers": 1, "pin_memory": True} if use_cuda else {}
+    data_dir = data_dir or "~/data"
+    with FileLock(os.path.expanduser("~/.horovod_lock")):
+        train_dataset = datasets.MNIST(
+            data_dir,
+            train=True,
+            download=True,
+            transform=transforms.Compose(
+                [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+            ),
+        )
+    # Horovod: use DistributedSampler to partition the training data.
+    train_sampler = torch.utils.data.distributed.DistributedSampler(
+        train_dataset, num_replicas=hvd.size(), rank=hvd.rank()
+    )
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=batch_size, sampler=train_sampler, **kwargs
+    )
+
+    model = Net()
+
+    # By default, Adasum doesn't need scaling up learning rate.
+    lr_scaler = hvd.size() if not use_adasum else 1
+
+    if use_cuda:
+        # Move model to GPU.
+        model.cuda()
+        # If using GPU Adasum allreduce, scale learning rate by local_size.
+        if use_adasum and hvd.nccl_built():
+            lr_scaler = hvd.local_size()
+
+    # Horovod: scale learning rate by lr_scaler.
+    optimizer = optim.SGD(model.parameters(), lr=lr * lr_scaler, momentum=momentum)
+
+    # Horovod: wrap optimizer with DistributedOptimizer.
+    optimizer = hvd.DistributedOptimizer(
+        optimizer,
+        named_parameters=model.named_parameters(),
+        op=hvd.Adasum if use_adasum else hvd.Average,
+    )
+
+    return model, optimizer, train_loader, train_sampler
+
+
+def train_epoch(
+    model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda
+):
+    loss = None
+    model.train()
+    # Horovod: set epoch to sampler for shuffling.
+    train_sampler.set_epoch(epoch)
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if use_cuda:
+            data, target = data.cuda(), target.cuda()
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % log_interval == 0:
+            # Horovod: use train_sampler to determine the number of
+            # examples in this worker's partition.
+            print(
+                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                    epoch,
+                    batch_idx * len(data),
+                    len(train_sampler),
+                    100.0 * batch_idx / len(train_loader),
+                    loss.item(),
+                )
+            )
+    return loss.item() if loss else None
+
+
+# Horovod function API.
+
+
+def train_func(config):
+    num_epochs = config.get("num_epochs", 10)
+    log_interval = config.get("log_interval", 10)
+    use_cuda = config.get("use_cuda", False)
+
+    model, optimizer, train_loader, train_sampler = setup(config)
+
+    for epoch in range(num_epochs):
+        loss = train_epoch(
+            model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda
+        )
+        train.report(dict(loss=loss))
+
+
+def main(num_workers, use_gpu, kwargs):
+    trainer = HorovodTrainer(
+        train_func,
+        train_loop_config=kwargs,
+        scaling_config=ScalingConfig(use_gpu=use_gpu, num_workers=num_workers),
+    )
+    results = trainer.fit()
+    print(results.metrics)
+
+
+# Horovod Class API.
+
+
+class HorovodTrainClass:
+    def __init__(self, config):
+        self.log_interval = config.get("log_interval", 10)
+        self.use_cuda = config.get("use_cuda", False)
+
+        if self.use_cuda:
+            torch.cuda.set_device(hvd.local_rank())
+
+        self.model, self.optimizer, self.train_loader, self.train_sampler = setup(
+            config
+        )
+
+    def train(self, epoch):
+        loss = train_epoch(
+            self.model,
+            self.optimizer,
+            self.train_sampler,
+            self.train_loader,
+            epoch,
+            self.log_interval,
+            self.use_cuda,
+        )
+        return loss
+
+
+if __name__ == "__main__":
+    # Training settings
+    parser = argparse.ArgumentParser(
+        description="PyTorch MNIST Example",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=64,
+        metavar="N",
+        help="input batch size for training (default: 64)",
+    )
+    parser.add_argument(
+        "--num-epochs",
+        type=int,
+        default=5,
+        metavar="N",
+        help="number of epochs to train (default: 10)",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=0.01,
+        metavar="LR",
+        help="learning rate (default: 0.01)",
+    )
+    parser.add_argument(
+        "--momentum",
+        type=float,
+        default=0.5,
+        metavar="M",
+        help="SGD momentum (default: 0.5)",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="enables CUDA training"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, metavar="S", help="random seed (default: 42)"
+    )
+    parser.add_argument(
+        "--log-interval",
+        type=int,
+        default=10,
+        metavar="N",
+        help="how many batches to wait before logging training status",
+    )
+    parser.add_argument(
+        "--use-adasum",
+        action="store_true",
+        default=False,
+        help="use adasum algorithm to do reduction",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=2,
+        help="Number of Ray workers to use for training.",
+    )
+    parser.add_argument(
+        "--data-dir",
+        help="location of the training dataset in the local filesystem ("
+        "will be downloaded if needed)",
+    )
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        default=None,
+        help="Address of Ray cluster.",
+    )
+
+    args = parser.parse_args()
+
+    if args.address:
+        ray.init(args.address)
+    else:
+        ray.init()
+
+    use_cuda = args.use_gpu if args.use_gpu is not None else False
+
+    kwargs = {
+        "data_dir": args.data_dir,
+        "seed": args.seed,
+        "use_cuda": use_cuda,
+        "batch_size": args.batch_size,
+        "use_adasum": args.use_adasum if args.use_adasum else False,
+        "lr": args.lr,
+        "momentum": args.momentum,
+        "num_epochs": args.num_epochs,
+        "log_interval": args.log_interval,
+    }
+
+    main(num_workers=args.num_workers, use_gpu=use_cuda, kwargs=kwargs)

python/ray/train/examples/horovod/horovod_pytorch_example.py

@@ -0,0 +1,270 @@
+import argparse
+import os
+import tempfile
+
+import horovod.torch as hvd
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.utils.data.distributed
+from filelock import FileLock
+from torchvision import datasets, transforms
+
+import ray.train.torch
+from ray import train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.horovod import HorovodTrainer
+
+
+def metric_average(val, name):
+    tensor = torch.tensor(val)
+    avg_tensor = hvd.allreduce(tensor, name=name)
+    return avg_tensor.item()
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x)
+
+
+def setup(config):
+    data_dir = config.get("data_dir", None)
+    seed = config.get("seed", 42)
+    batch_size = config.get("batch_size", 64)
+    use_adasum = config.get("use_adasum", False)
+    lr = config.get("lr", 0.01)
+    momentum = config.get("momentum", 0.5)
+    use_cuda = config.get("use_cuda", False)
+
+    # Horovod: initialize library.
+    hvd.init()
+    torch.manual_seed(seed)
+
+    if use_cuda:
+        # Horovod: pin GPU to local rank.
+        torch.cuda.set_device(hvd.local_rank())
+        torch.cuda.manual_seed(seed)
+
+    # Horovod: limit # of CPU threads to be used per worker.
+    torch.set_num_threads(1)
+
+    kwargs = {"pin_memory": True} if use_cuda else {}
+    data_dir = data_dir or "~/data"
+    with FileLock(os.path.expanduser("~/.horovod_lock")):
+        train_dataset = datasets.MNIST(
+            data_dir,
+            train=True,
+            download=True,
+            transform=transforms.Compose(
+                [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+            ),
+        )
+    # Horovod: use DistributedSampler to partition the training data.
+    train_sampler = torch.utils.data.distributed.DistributedSampler(
+        train_dataset, num_replicas=hvd.size(), rank=hvd.rank()
+    )
+    # Note, don't set `num_workers` in DataLoader (not even 1),
+    # as that will separately start multiple processes (each corresponding to 1 worker)
+    # to load the data. This is known to cause issues with Ray.
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=batch_size, sampler=train_sampler, **kwargs
+    )
+
+    model = Net()
+
+    # By default, Adasum doesn't need scaling up learning rate.
+    lr_scaler = hvd.size() if not use_adasum else 1
+
+    if use_cuda:
+        # Move model to GPU.
+        model.cuda()
+        # If using GPU Adasum allreduce, scale learning rate by local_size.
+        if use_adasum and hvd.nccl_built():
+            lr_scaler = hvd.local_size()
+
+    # Horovod: scale learning rate by lr_scaler.
+    optimizer = optim.SGD(model.parameters(), lr=lr * lr_scaler, momentum=momentum)
+
+    # Horovod: wrap optimizer with DistributedOptimizer.
+    optimizer = hvd.DistributedOptimizer(
+        optimizer,
+        named_parameters=model.named_parameters(),
+        op=hvd.Adasum if use_adasum else hvd.Average,
+    )
+
+    return model, optimizer, train_loader, train_sampler
+
+
+def train_epoch(
+    model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda
+):
+    loss = None
+    model.train()
+    # Horovod: set epoch to sampler for shuffling.
+    train_sampler.set_epoch(epoch)
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if use_cuda:
+            data, target = data.cuda(), target.cuda()
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % log_interval == 0:
+            # Horovod: use train_sampler to determine the number of
+            # examples in this worker's partition.
+            print(
+                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                    epoch,
+                    batch_idx * len(data),
+                    len(train_sampler),
+                    100.0 * batch_idx / len(train_loader),
+                    loss.item(),
+                )
+            )
+    return loss.item() if loss else None
+
+
+def train_func(config):
+    num_epochs = config.get("num_epochs", 10)
+    log_interval = config.get("log_interval", 10)
+    use_cuda = config.get("use_cuda", False)
+
+    model, optimizer, train_loader, train_sampler = setup(config)
+
+    results = []
+    for epoch in range(num_epochs):
+        loss = train_epoch(
+            model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda
+        )
+        results.append(loss)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(model.state_dict(), os.path.join(tmpdir, "model.pt"))
+            train.report({"loss": loss}, checkpoint=Checkpoint.from_directory(tmpdir))
+
+    # Only used for testing.
+    return results
+
+
+def main(num_workers, use_gpu, kwargs):
+    trainer = HorovodTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config={
+            "num_epochs": kwargs["num_epochs"],
+            "log_interval": kwargs["log_interval"],
+            "use_cuda": kwargs["use_cuda"],
+        },
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+    )
+    result = trainer.fit()
+    print(result)
+
+
+if __name__ == "__main__":
+    # Training settings
+    parser = argparse.ArgumentParser(
+        description="PyTorch MNIST Example",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=64,
+        metavar="N",
+        help="input batch size for training (default: 64)",
+    )
+    parser.add_argument(
+        "--num-epochs",
+        type=int,
+        default=5,
+        metavar="N",
+        help="number of epochs to train (default: 10)",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=0.01,
+        metavar="LR",
+        help="learning rate (default: 0.01)",
+    )
+    parser.add_argument(
+        "--momentum",
+        type=float,
+        default=0.5,
+        metavar="M",
+        help="SGD momentum (default: 0.5)",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="enables CUDA training"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, metavar="S", help="random seed (default: 42)"
+    )
+    parser.add_argument(
+        "--log-interval",
+        type=int,
+        default=10,
+        metavar="N",
+        help="how many batches to wait before logging training status",
+    )
+    parser.add_argument(
+        "--use-adasum",
+        action="store_true",
+        default=False,
+        help="use adasum algorithm to do reduction",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=2,
+        help="Number of Ray workers to use for training.",
+    )
+    parser.add_argument(
+        "--data-dir",
+        help="location of the training dataset in the local filesystem ("
+        "will be downloaded if needed)",
+    )
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        default=None,
+        help="Address of Ray cluster.",
+    )
+
+    args = parser.parse_args()
+
+    if args.address:
+        ray.init(args.address)
+    else:
+        ray.init()
+
+    use_cuda = args.use_gpu if args.use_gpu is not None else False
+
+    kwargs = {
+        "data_dir": args.data_dir,
+        "seed": args.seed,
+        "use_cuda": use_cuda,
+        "batch_size": args.batch_size,
+        "use_adasum": args.use_adasum if args.use_adasum else False,
+        "lr": args.lr,
+        "momentum": args.momentum,
+        "num_epochs": args.num_epochs,
+        "log_interval": args.log_interval,
+    }
+
+    main(num_workers=args.num_workers, use_gpu=use_cuda, kwargs=kwargs)

python/ray/train/examples/horovod/horovod_tune_example.py

@@ -0,0 +1,139 @@
+import time
+
+import numpy as np
+import torch
+
+import ray
+import ray.train.torch
+from ray import train, tune
+from ray.train import ScalingConfig
+from ray.train.horovod import HorovodTrainer
+from ray.tune.tune_config import TuneConfig
+from ray.tune.tuner import Tuner
+
+
+def sq(x):
+    m2 = 1.0
+    m1 = -20.0
+    m0 = 50.0
+    return m2 * x * x + m1 * x + m0
+
+
+def qu(x):
+    m3 = 10.0
+    m2 = 5.0
+    m1 = -20.0
+    m0 = -5.0
+    return m3 * x * x * x + m2 * x * x + m1 * x + m0
+
+
+class Net(torch.nn.Module):
+    def __init__(self, mode="sq"):
+        super(Net, self).__init__()
+
+        if mode == "square":
+            self.mode = 0
+            self.param = torch.nn.Parameter(torch.FloatTensor([1.0, -1.0]))
+        else:
+            self.mode = 1
+            self.param = torch.nn.Parameter(torch.FloatTensor([1.0, -1.0, 1.0]))
+
+    def forward(self, x):
+        if ~self.mode:
+            return x * x + self.param[0] * x + self.param[1]
+        else:
+            return_val = 10 * x * x * x
+            return_val += self.param[0] * x * x
+            return_val += self.param[1] * x + self.param[2]
+            return return_val
+
+
+def train_loop_per_worker(config):
+    import horovod.torch as hvd
+    import torch
+
+    hvd.init()
+    device = ray.train.torch.get_device()
+    mode = config["mode"]
+    net = Net(mode).to(device)
+    optimizer = torch.optim.SGD(
+        net.parameters(),
+        lr=config["lr"],
+    )
+    optimizer = hvd.DistributedOptimizer(optimizer)
+
+    num_steps = 5
+    print(hvd.size())
+    np.random.seed(1 + hvd.rank())
+    torch.manual_seed(1234)
+    # To ensure consistent initialization across workers,
+    hvd.broadcast_parameters(net.state_dict(), root_rank=0)
+    hvd.broadcast_optimizer_state(optimizer, root_rank=0)
+
+    start = time.time()
+    x_max = config["x_max"]
+    for step in range(1, num_steps + 1):
+        features = torch.Tensor(np.random.rand(1) * 2 * x_max - x_max).to(device)
+        if mode == "square":
+            labels = sq(features)
+        else:
+            labels = qu(features)
+        optimizer.zero_grad()
+        outputs = net(features)
+        loss = torch.nn.MSELoss()(outputs, labels)
+        loss.backward()
+
+        optimizer.step()
+        time.sleep(0.1)
+        train.report(dict(loss=loss.item()))
+    total = time.time() - start
+    print(f"Took {total:0.3f} s. Avg: {total / num_steps:0.3f} s.")
+
+
+def tune_horovod(num_workers, num_samples, use_gpu, mode="square", x_max=1.0):
+    horovod_trainer = HorovodTrainer(
+        train_loop_per_worker=train_loop_per_worker,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+        train_loop_config={"mode": mode, "x_max": x_max},
+    )
+
+    tuner = Tuner(
+        horovod_trainer,
+        param_space={"train_loop_config": {"lr": tune.uniform(0.1, 1)}},
+        tune_config=TuneConfig(mode="min", metric="loss", num_samples=num_samples),
+        _tuner_kwargs={"fail_fast": True},
+    )
+
+    result_grid = tuner.fit()
+
+    print("Best hyperparameters found were: ", result_grid.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--mode", type=str, default="square", choices=["square", "cubic"]
+    )
+    parser.add_argument(
+        "--learning_rate", type=float, default=0.1, dest="learning_rate"
+    )
+    parser.add_argument("--x_max", type=float, default=1.0, dest="x_max")
+    parser.add_argument("--gpu", action="store_true")
+    parser.add_argument(
+        "--smoke-test", action="store_true", help=("Finish quickly for testing.")
+    )
+    parser.add_argument("--num-workers", type=int, default=2)
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        ray.init(num_cpus=3)
+
+    tune_horovod(
+        num_workers=args.num_workers,
+        num_samples=2 if args.smoke_test else 10,
+        use_gpu=args.gpu,
+        mode=args.mode,
+        x_max=args.x_max,
+    )

python/ray/train/examples/pytorch/torch_data_prefetch_benchmark/auto_pipeline_for_host_to_device_data_transfer.py

@@ -0,0 +1,161 @@
+# The PyTorch data transfer benchmark script.
+import argparse
+import warnings
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import ray.train as train
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+class Net(nn.Module):
+    def __init__(self, in_d, hidden):
+        # output dim = 1
+        super(Net, self).__init__()
+        dims = [in_d] + hidden + [1]
+        self.layers = nn.ModuleList(
+            [nn.Linear(dims[i - 1], dims[i]) for i in range(len(dims))]
+        )
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class BenchmarkDataset(torch.utils.data.Dataset):
+    """Create a naive dataset for the benchmark"""
+
+    def __init__(self, dim, size=1000):
+        self.x = torch.from_numpy(np.random.normal(size=(size, dim))).float()
+        self.y = torch.from_numpy(np.random.normal(size=(size, 1))).float()
+        self.size = size
+
+    def __getitem__(self, index):
+        return self.x[index, None], self.y[index, None]
+
+    def __len__(self):
+        return self.size
+
+
+def train_epoch(epoch, dataloader, model, loss_fn, optimizer):
+    if train.get_context().get_world_size() > 1:
+        dataloader.sampler.set_epoch(epoch)
+
+    for X, y in dataloader:
+        # Compute prediction error
+        pred = model(X)
+        loss = loss_fn(pred, y)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+
+def train_func(config):
+    data_size = config.get("data_size", 4096 * 50)
+    batch_size = config.get("batch_size", 4096)
+    hidden_size = config.get("hidden_size", 1)
+    use_auto_transfer = config.get("use_auto_transfer", False)
+    lr = config.get("lr", 1e-2)
+    epochs = config.get("epochs", 10)
+
+    train_dataset = BenchmarkDataset(4096, size=data_size)
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=batch_size, shuffle=True
+    )
+
+    train_loader = train.torch.prepare_data_loader(
+        data_loader=train_loader, move_to_device=True, auto_transfer=use_auto_transfer
+    )
+
+    model = Net(in_d=4096, hidden=[4096] * hidden_size)
+    model = train.torch.prepare_model(model)
+
+    loss_fn = nn.MSELoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+
+    choice = "with" if use_auto_transfer else "without"
+    print(f"Starting the torch data prefetch benchmark {choice} auto pipeline...")
+
+    torch.cuda.synchronize()
+    start.record()
+    for epoch in range(epochs):
+        train_epoch(epoch, train_loader, model, loss_fn, optimizer)
+    end.record()
+    torch.cuda.synchronize()
+
+    print(
+        f"Finished the torch data prefetch benchmark {choice} "
+        f"auto pipeline: {start.elapsed_time(end)} ms."
+    )
+
+    return "Experiment done."
+
+
+def train_linear(num_workers=1, num_hidden_layers=1, use_auto_transfer=True, epochs=3):
+    config = {
+        "lr": 1e-2,
+        "hidden_size": num_hidden_layers,
+        "batch_size": 4096,
+        "epochs": epochs,
+        "use_auto_transfer": use_auto_transfer,
+    }
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(use_gpu=True, num_workers=num_workers),
+    )
+    results = trainer.fit()
+
+    print(results.metrics)
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=1, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--num_hidden_layers",
+        type=int,
+        default=1,
+        help="Number of epochs to train for.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    import ray
+
+    ray.init(address=args.address)
+
+    if not torch.cuda.is_available():
+        warnings.warn("GPU is not available. Skip the test using auto pipeline.")
+    else:
+        train_linear(
+            num_workers=1,
+            num_hidden_layers=args.num_hidden_layers,
+            use_auto_transfer=True,
+            epochs=args.epochs,
+        )
+
+    torch.cuda.empty_cache()
+    train_linear(
+        num_workers=1,
+        num_hidden_layers=args.num_hidden_layers,
+        use_auto_transfer=False,
+        epochs=args.epochs,
+    )
+
+    ray.shutdown()

python/ray/train/examples/pytorch/torch_fashion_mnist_example.py

@@ -0,0 +1,154 @@
+import os
+from typing import Dict
+
+import torch
+from filelock import FileLock
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from torchvision.transforms import Normalize, ToTensor
+from tqdm import tqdm
+
+import ray.train
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def get_dataloaders(batch_size):
+    # Transform to normalize the input images
+    transform = transforms.Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
+
+    with FileLock(os.path.expanduser("~/data.lock")):
+        # Download training data from open datasets
+        training_data = datasets.FashionMNIST(
+            root="~/data",
+            train=True,
+            download=True,
+            transform=transform,
+        )
+
+        # Download test data from open datasets
+        test_data = datasets.FashionMNIST(
+            root="~/data",
+            train=False,
+            download=True,
+            transform=transform,
+        )
+
+    # Create data loaders
+    train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
+    test_dataloader = DataLoader(test_data, batch_size=batch_size)
+
+    return train_dataloader, test_dataloader
+
+
+# Model Definition
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super(NeuralNetwork, self).__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(28 * 28, 512),
+            nn.ReLU(),
+            nn.Dropout(0.25),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Dropout(0.25),
+            nn.Linear(512, 10),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        x = self.flatten(x)
+        logits = self.linear_relu_stack(x)
+        return logits
+
+
+def train_func_per_worker(config: Dict):
+    ray.train.torch.enable_reproducibility()
+
+    lr = config["lr"]
+    epochs = config["epochs"]
+    batch_size = config["batch_size_per_worker"]
+
+    # Get dataloaders inside the worker training function
+    train_dataloader, test_dataloader = get_dataloaders(batch_size=batch_size)
+
+    # [1] Prepare Dataloader for distributed training
+    # Shard the datasets among workers and move batches to the correct device
+    # =======================================================================
+    train_dataloader = ray.train.torch.prepare_data_loader(train_dataloader)
+    test_dataloader = ray.train.torch.prepare_data_loader(test_dataloader)
+
+    model = NeuralNetwork()
+
+    # [2] Prepare and wrap your model with DistributedDataParallel
+    # Move the model to the correct GPU/CPU device
+    # ============================================================
+    model = ray.train.torch.prepare_model(model)
+
+    loss_fn = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
+
+    # Model training loop
+    for epoch in range(epochs):
+        if ray.train.get_context().get_world_size() > 1:
+            # Required for the distributed sampler to shuffle properly across epochs.
+            train_dataloader.sampler.set_epoch(epoch)
+
+        model.train()
+        for X, y in tqdm(train_dataloader, desc=f"Train Epoch {epoch}"):
+            pred = model(X)
+            loss = loss_fn(pred, y)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        model.eval()
+        test_loss, num_correct, num_total = 0, 0, 0
+        with torch.no_grad():
+            for X, y in tqdm(test_dataloader, desc=f"Test Epoch {epoch}"):
+                pred = model(X)
+                loss = loss_fn(pred, y)
+
+                test_loss += loss.item()
+                num_total += y.shape[0]
+                num_correct += (pred.argmax(1) == y).sum().item()
+
+        test_loss /= len(test_dataloader)
+        accuracy = num_correct / num_total
+
+        # [3] Report metrics to Ray Train
+        # ===============================
+        ray.train.report(metrics={"loss": test_loss, "accuracy": accuracy})
+
+
+def train_fashion_mnist(num_workers=2, use_gpu=False):
+    global_batch_size = 32
+
+    train_config = {
+        "lr": 1e-3,
+        "epochs": 10,
+        "batch_size_per_worker": global_batch_size // num_workers,
+    }
+
+    # Configure computation resources
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
+
+    # Initialize a Ray TorchTrainer
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func_per_worker,
+        train_loop_config=train_config,
+        scaling_config=scaling_config,
+    )
+
+    # [4] Start distributed training
+    # Run `train_func_per_worker` on all workers
+    # =============================================
+    result = trainer.fit()
+    print(f"Training result: {result}")
+
+
+if __name__ == "__main__":
+    train_fashion_mnist(num_workers=4, use_gpu=True)

python/ray/train/examples/pytorch/torch_linear_example.py

@@ -0,0 +1,147 @@
+import argparse
+import os
+import tempfile
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import ray.train as train
+from ray.train import Checkpoint, RunConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+class LinearDataset(torch.utils.data.Dataset):
+    """y = a * x + b"""
+
+    def __init__(self, a, b, size=1000):
+        x = np.arange(0, 10, 10 / size, dtype=np.float32)
+        self.x = torch.from_numpy(x)
+        self.y = torch.from_numpy(a * x + b)
+
+    def __getitem__(self, index):
+        return self.x[index, None], self.y[index, None]
+
+    def __len__(self):
+        return len(self.x)
+
+
+def train_epoch(epoch, dataloader, model, loss_fn, optimizer):
+    if train.get_context().get_world_size() > 1:
+        dataloader.sampler.set_epoch(epoch)
+
+    for X, y in dataloader:
+        # Compute prediction error
+        pred = model(X)
+        loss = loss_fn(pred, y)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+
+def validate_epoch(dataloader, model, loss_fn):
+    num_batches = len(dataloader)
+    model.eval()
+    loss = 0
+    with torch.no_grad():
+        for X, y in dataloader:
+            pred = model(X)
+            loss += loss_fn(pred, y).item()
+    loss /= num_batches
+    import copy
+
+    model_copy = copy.deepcopy(model)
+    return model_copy.cpu().state_dict(), loss
+
+
+def train_func(config):
+    data_size = config.get("data_size", 1000)
+    val_size = config.get("val_size", 400)
+    batch_size = config.get("batch_size", 32)
+    hidden_size = config.get("hidden_size", 1)
+    lr = config.get("lr", 1e-2)
+    epochs = config.get("epochs", 3)
+
+    train_dataset = LinearDataset(2, 5, size=data_size)
+    val_dataset = LinearDataset(2, 5, size=val_size)
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)
+    validation_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size)
+
+    train_loader = train.torch.prepare_data_loader(train_loader)
+    validation_loader = train.torch.prepare_data_loader(validation_loader)
+
+    model = nn.Linear(1, hidden_size)
+    model = train.torch.prepare_model(model)
+
+    loss_fn = nn.MSELoss()
+
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+
+    results = []
+    for epoch in range(epochs):
+        train_epoch(epoch, train_loader, model, loss_fn, optimizer)
+        state_dict, loss = validate_epoch(validation_loader, model, loss_fn)
+        result = dict(loss=loss)
+        results.append(result)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(state_dict, os.path.join(tmpdir, "model.pt"))
+            train.report(result, checkpoint=Checkpoint.from_directory(tmpdir))
+
+    return results
+
+
+def train_linear(num_workers=2, use_gpu=False, epochs=3, storage_path=None):
+    config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+        run_config=RunConfig(storage_path=storage_path),
+    )
+    result = trainer.fit()
+
+    print(result.metrics)
+    return result.metrics
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", help="Whether to use GPU for training."
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    import ray
+
+    if args.smoke_test:
+        # 2 workers + 1 for trainer.
+        ray.init(num_cpus=3)
+        train_linear()
+    else:
+        ray.init(address=args.address)
+        train_linear(
+            num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
+        )

python/ray/train/examples/pytorch/torch_quick_start.py

@@ -0,0 +1,110 @@
+# ruff: noqa
+# fmt: off
+# isort: skip_file
+
+# __torch_setup_begin__
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision.transforms import ToTensor
+
+def get_dataset():
+    return datasets.FashionMNIST(
+        root="/tmp/data",
+        train=True,
+        download=True,
+        transform=ToTensor(),
+    )
+
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(28 * 28, 512),
+            nn.ReLU(),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Linear(512, 10),
+        )
+
+    def forward(self, inputs):
+        inputs = self.flatten(inputs)
+        logits = self.linear_relu_stack(inputs)
+        return logits
+# __torch_setup_end__
+
+# __torch_single_begin__
+def train_func():
+    num_epochs = 3
+    batch_size = 64
+
+    dataset = get_dataset()
+    dataloader = DataLoader(dataset, batch_size=batch_size)
+
+    model = NeuralNetwork()
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+    for epoch in range(num_epochs):
+        for inputs, labels in dataloader:
+            optimizer.zero_grad()
+            pred = model(inputs)
+            loss = criterion(pred, labels)
+            loss.backward()
+            optimizer.step()
+        print(f"epoch: {epoch}, loss: {loss.item()}")
+# __torch_single_end__
+
+# __torch_distributed_begin__
+import ray.train.torch
+
+def train_func_distributed():
+    num_epochs = 3
+    batch_size = 64
+
+    dataset = get_dataset()
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+    dataloader = ray.train.torch.prepare_data_loader(dataloader)
+
+    model = NeuralNetwork()
+    model = ray.train.torch.prepare_model(model)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+    for epoch in range(num_epochs):
+        if ray.train.get_context().get_world_size() > 1:
+            dataloader.sampler.set_epoch(epoch)
+
+        for inputs, labels in dataloader:
+            optimizer.zero_grad()
+            pred = model(inputs)
+            loss = criterion(pred, labels)
+            loss.backward()
+            optimizer.step()
+        print(f"epoch: {epoch}, loss: {loss.item()}")
+# __torch_distributed_end__
+
+
+if __name__ == "__main__":
+    # __torch_single_run_begin__
+    train_func()
+    # __torch_single_run_end__
+
+    # __torch_trainer_begin__
+    from ray.train.torch import TorchTrainer
+    from ray.train import ScalingConfig
+
+    # For GPU Training, set `use_gpu` to True.
+    use_gpu = False
+
+    trainer = TorchTrainer(
+        train_func_distributed,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=use_gpu)
+    )
+
+    results = trainer.fit()
+    # __torch_trainer_end__

python/ray/train/examples/pytorch/torch_regression_example.py

@@ -0,0 +1,159 @@
+import argparse
+import os
+import tempfile
+from typing import Tuple
+
+import pandas as pd
+import torch
+import torch.nn as nn
+
+import ray
+import ray.train as train
+from ray.data import Dataset
+from ray.train import Checkpoint, DataConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def get_datasets(split: float = 0.7) -> Tuple[Dataset]:
+    dataset = ray.data.read_csv("s3://anonymous@air-example-data/regression.csv")
+
+    def combine_x(batch):
+        return pd.DataFrame(
+            {
+                "x": batch[[f"x{i:03d}" for i in range(100)]].values.tolist(),
+                "y": batch["y"],
+            }
+        )
+
+    dataset = dataset.map_batches(combine_x, batch_format="pandas")
+    train_dataset, validation_dataset = dataset.repartition(
+        num_blocks=4
+    ).train_test_split(split, shuffle=True)
+    return train_dataset, validation_dataset
+
+
+def train_epoch(iterable_dataset, model, loss_fn, optimizer, device):
+    model.train()
+    for X, y in iterable_dataset:
+        X = X.to(device)
+        y = y.to(device)
+
+        # Compute prediction error
+        pred = model(X)
+        loss = loss_fn(pred, y)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+
+def validate_epoch(iterable_dataset, model, loss_fn, device):
+    num_batches = 0
+    model.eval()
+    loss = 0
+    with torch.no_grad():
+        for X, y in iterable_dataset:
+            X = X.to(device)
+            y = y.to(device)
+            num_batches += 1
+            pred = model(X)
+            loss += loss_fn(pred, y).item()
+    loss /= num_batches
+    result = {"loss": loss}
+    return result
+
+
+def train_func(config):
+    batch_size = config.get("batch_size", 32)
+    hidden_size = config.get("hidden_size", 10)
+    lr = config.get("lr", 1e-2)
+    epochs = config.get("epochs", 3)
+
+    train_dataset_shard = train.get_dataset_shard("train")
+    validation_dataset = train.get_dataset_shard("validation")
+
+    model = nn.Sequential(
+        nn.Linear(100, hidden_size), nn.ReLU(), nn.Linear(hidden_size, 1)
+    )
+    model = train.torch.prepare_model(model)
+
+    loss_fn = nn.L1Loss()
+
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+
+    results = []
+
+    def create_torch_iterator(shard):
+        iterator = shard.iter_torch_batches(batch_size=batch_size)
+        for batch in iterator:
+            yield batch["x"].float(), batch["y"].float()
+
+    for _ in range(epochs):
+        train_torch_dataset = create_torch_iterator(train_dataset_shard)
+        validation_torch_dataset = create_torch_iterator(validation_dataset)
+
+        device = train.torch.get_device()
+
+        train_epoch(train_torch_dataset, model, loss_fn, optimizer, device)
+        if train.get_context().get_world_rank() == 0:
+            result = validate_epoch(validation_torch_dataset, model, loss_fn, device)
+        else:
+            result = {}
+        results.append(result)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(model.module.state_dict(), os.path.join(tmpdir, "model.pt"))
+            train.report(result, checkpoint=Checkpoint.from_directory(tmpdir))
+
+    return results
+
+
+def train_regression(num_workers=2, use_gpu=False):
+    train_dataset, val_dataset = get_datasets()
+    config = {"lr": 1e-2, "hidden_size": 20, "batch_size": 4, "epochs": 3}
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+        datasets={"train": train_dataset, "validation": val_dataset},
+        dataset_config=DataConfig(datasets_to_split=["train"]),
+    )
+
+    result = trainer.fit()
+    return result
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Use GPU for training."
+    )
+
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        ray.init(num_cpus=4)
+        result = train_regression()
+    else:
+        ray.init(address=args.address)
+        result = train_regression(num_workers=args.num_workers, use_gpu=args.use_gpu)
+    print(result)

python/ray/train/examples/pytorch_geometric/distributed_sage_example.py

@@ -0,0 +1,228 @@
+# Adapted from https://github.com/pyg-team/pytorch_geometric/blob/2.1.0
+# /examples/multi_gpu/distributed_sampling.py
+
+import argparse
+import os
+
+import torch
+import torch.nn.functional as F
+from filelock import FileLock
+from torch_geometric.datasets import FakeDataset, Reddit
+from torch_geometric.loader import NeighborSampler
+from torch_geometric.nn import SAGEConv
+from torch_geometric.transforms import RandomNodeSplit
+
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+class SAGE(torch.nn.Module):
+    def __init__(self, in_channels, hidden_channels, out_channels, num_layers=2):
+        super().__init__()
+        self.num_layers = num_layers
+
+        self.convs = torch.nn.ModuleList()
+        self.convs.append(SAGEConv(in_channels, hidden_channels))
+        for _ in range(self.num_layers - 2):
+            self.convs.append(SAGEConv(hidden_channels, hidden_channels))
+        self.convs.append(SAGEConv(hidden_channels, out_channels))
+
+    def forward(self, x, adjs):
+        for i, (edge_index, _, size) in enumerate(adjs):
+            x_target = x[: size[1]]  # Target nodes are always placed first.
+            x = self.convs[i]((x, x_target), edge_index)
+            if i != self.num_layers - 1:
+                x = F.relu(x)
+                x = F.dropout(x, p=0.5, training=self.training)
+        return x.log_softmax(dim=-1)
+
+    @torch.no_grad()
+    def test(self, x_all, subgraph_loader):
+        for i in range(self.num_layers):
+            xs = []
+            for batch_size, n_id, adj in subgraph_loader:
+                edge_index, _, size = adj
+                x = x_all[n_id.to(x_all.device)].to(train.torch.get_device())
+                x_target = x[: size[1]]
+                x = self.convs[i]((x, x_target), edge_index)
+                if i != self.num_layers - 1:
+                    x = F.relu(x)
+                xs.append(x.cpu())
+
+            x_all = torch.cat(xs, dim=0)
+
+        return x_all
+
+
+def train_loop_per_worker(train_loop_config):
+    dataset = train_loop_config["dataset_fn"]()
+    batch_size = train_loop_config["batch_size"]
+    num_epochs = train_loop_config["num_epochs"]
+
+    data = dataset[0]
+    train_idx = data.train_mask.nonzero(as_tuple=False).view(-1)
+    train_idx = train_idx.split(
+        train_idx.size(0) // train.get_context().get_world_size()
+    )[train.get_context().get_world_rank()]
+
+    train_loader = NeighborSampler(
+        data.edge_index,
+        node_idx=train_idx,
+        sizes=[25, 10],
+        batch_size=batch_size,
+        shuffle=True,
+    )
+
+    # Disable distributed sampler since the train_loader has already been split above.
+    train_loader = train.torch.prepare_data_loader(train_loader, add_dist_sampler=False)
+
+    # Do validation on rank 0 worker only.
+    if train.get_context().get_world_rank() == 0:
+        subgraph_loader = NeighborSampler(
+            data.edge_index, node_idx=None, sizes=[-1], batch_size=2048, shuffle=False
+        )
+        subgraph_loader = train.torch.prepare_data_loader(
+            subgraph_loader, add_dist_sampler=False
+        )
+
+    model = SAGE(dataset.num_features, 256, dataset.num_classes)
+    model = train.torch.prepare_model(model)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+    x, y = data.x.to(train.torch.get_device()), data.y.to(train.torch.get_device())
+
+    for epoch in range(num_epochs):
+        model.train()
+
+        # ``batch_size`` is the number of samples in the current batch.
+        # ``n_id`` are the ids of all the nodes used in the computation. This is
+        # needed to pull in the necessary features just for the current batch that is
+        # being trained on.
+        # ``adjs`` is a list of 3 element tuple consisting of ``(edge_index, e_id,
+        # size)`` for each sample in the batch, where ``edge_index``represent the
+        # edges of the sampled subgraph, ``e_id`` are the ids of the edges in the
+        # sample, and ``size`` holds the shape of the subgraph.
+        # See ``torch_geometric.loader.neighbor_sampler.NeighborSampler`` for more info.
+        for batch_size, n_id, adjs in train_loader:
+            optimizer.zero_grad()
+            out = model(x[n_id], adjs)
+            loss = F.nll_loss(out, y[n_id[:batch_size]])
+            loss.backward()
+            optimizer.step()
+
+        if train.get_context().get_world_rank() == 0:
+            print(f"Epoch: {epoch:03d}, Loss: {loss:.4f}")
+
+        train_accuracy = validation_accuracy = test_accuracy = None
+
+        # Do validation on rank 0 worker only.
+        if train.get_context().get_world_rank() == 0:
+            model.eval()
+            with torch.no_grad():
+                out = model.module.test(x, subgraph_loader)
+            res = out.argmax(dim=-1) == data.y
+            train_accuracy = int(res[data.train_mask].sum()) / int(
+                data.train_mask.sum()
+            )
+            validation_accuracy = int(res[data.val_mask].sum()) / int(
+                data.val_mask.sum()
+            )
+            test_accuracy = int(res[data.test_mask].sum()) / int(data.test_mask.sum())
+
+        train.report(
+            dict(
+                train_accuracy=train_accuracy,
+                validation_accuracy=validation_accuracy,
+                test_accuracy=test_accuracy,
+            )
+        )
+
+
+def gen_fake_dataset():
+    """Returns a function to be called on each worker that returns a Fake Dataset."""
+
+    # For fake dataset, since the dataset is randomized, we create it once on the
+    # driver, and then send the same dataset to all the training workers.
+    # Use 10% of nodes for validation and 10% for testing.
+    fake_dataset = FakeDataset(transform=RandomNodeSplit(num_val=0.1, num_test=0.1))
+
+    def gen_dataset():
+        return fake_dataset
+
+    return gen_dataset
+
+
+def gen_reddit_dataset():
+    """Returns a function to be called on each worker that returns Reddit Dataset."""
+
+    # For Reddit dataset, we have to download the data on each node, so we create the
+    # dataset on each training worker.
+    with FileLock(os.path.expanduser("~/.reddit_dataset_lock")):
+        dataset = Reddit("./data/Reddit")
+    return dataset
+
+
+def train_gnn(
+    num_workers=2, use_gpu=False, epochs=3, global_batch_size=32, dataset="reddit"
+):
+    per_worker_batch_size = global_batch_size // num_workers
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop_per_worker,
+        train_loop_config={
+            "num_epochs": epochs,
+            "batch_size": per_worker_batch_size,
+            "dataset_fn": gen_reddit_dataset
+            if dataset == "reddit"
+            else gen_fake_dataset(),
+        },
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+    )
+    result = trainer.fit()
+    print(result.metrics)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", help="Whether to use GPU for training."
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--global-batch-size",
+        "-b",
+        type=int,
+        default=32,
+        help="Global batch size to use for training.",
+    )
+    parser.add_argument(
+        "--dataset",
+        "-d",
+        type=str,
+        choices=["reddit", "fake"],
+        default="reddit",
+        help="The dataset to use. Either 'reddit' or 'fake' Defaults to 'reddit'.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    train_gnn(
+        num_workers=args.num_workers,
+        use_gpu=args.use_gpu,
+        epochs=args.epochs,
+        global_batch_size=args.global_batch_size,
+        dataset=args.dataset,
+    )

python/ray/train/examples/tf/tensorflow_autoencoder_example.py

@@ -0,0 +1,177 @@
+# This example showcases how to use Tensorflow with Ray Train.
+# Original code:
+# https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import argparse
+
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+import tensorflow_datasets as tfds
+
+import ray
+from ray import train
+from ray.air.integrations.keras import ReportCheckpointCallback
+from ray.data.datasource import SimpleTensorFlowDatasource
+from ray.data.extensions import TensorArray
+from ray.train import Result, ScalingConfig
+from ray.train.tensorflow import TensorflowTrainer, prepare_dataset_shard
+
+
+def get_dataset(split_type="train"):
+    def dataset_factory():
+        return tfds.load("mnist", split=[split_type], as_supervised=True)[0].take(128)
+
+    dataset = ray.data.read_datasource(
+        SimpleTensorFlowDatasource(), dataset_factory=dataset_factory
+    )
+
+    def normalize_images(x):
+        x = np.float32(x.numpy()) / 255.0
+        x = np.reshape(x, (-1,))
+        return x
+
+    def preprocess_dataset(batch):
+        return [
+            (normalize_images(image), normalize_images(image)) for image, _ in batch
+        ]
+
+    dataset = dataset.map_batches(preprocess_dataset)
+
+    def convert_batch_to_pandas(batch):
+
+        images = [TensorArray(image) for image, _ in batch]
+        # because we did autoencoder here
+        df = pd.DataFrame({"image": images, "label": images})
+        return df
+
+    dataset = dataset.map_batches(convert_batch_to_pandas)
+    return dataset
+
+
+def build_autoencoder_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=(784,)),
+            # encoder
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(64, activation="relu"),
+            tf.keras.layers.Dense(32, activation="relu"),
+            # decoder
+            tf.keras.layers.Dense(64, activation="relu"),
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(784, activation="sigmoid"),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+
+    per_worker_batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+
+    dataset_shard = train.get_dataset_shard("train")
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_autoencoder_model()
+        learning_rate = config.get("lr", 0.001)
+        multi_worker_model.compile(
+            loss=tf.keras.losses.BinaryCrossentropy(),
+            optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
+            metrics=[
+                "binary_crossentropy",
+            ],
+        )
+
+    def to_tf_dataset(dataset, batch_size):
+        def to_tensor_iterator():
+            for batch in dataset.iter_tf_batches(
+                batch_size=batch_size, dtypes=tf.float32
+            ):
+                yield batch["image"], batch["label"]
+
+        output_signature = (
+            tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
+            tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
+        )
+        tf_dataset = tf.data.Dataset.from_generator(
+            to_tensor_iterator, output_signature=output_signature
+        )
+        return prepare_dataset_shard(tf_dataset)
+
+    results = []
+    for epoch in range(epochs):
+        tf_dataset = to_tf_dataset(
+            dataset=dataset_shard,
+            batch_size=per_worker_batch_size,
+        )
+        history = multi_worker_model.fit(
+            tf_dataset, callbacks=[ReportCheckpointCallback()]
+        )
+        results.append(history.history)
+    return results
+
+
+def train_tensorflow_mnist(
+    num_workers: int = 2, use_gpu: bool = False, epochs: int = 4
+) -> Result:
+    train_dataset = get_dataset(split_type="train")
+    config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        datasets={"train": train_dataset},
+        scaling_config=scaling_config,
+    )
+
+    results = trainer.fit()
+    print(results.metrics)
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        num_gpus = args.num_workers if args.use_gpu else 0
+        ray.init(num_cpus=4, num_gpus=num_gpus)
+        result = train_tensorflow_mnist(num_workers=2, use_gpu=args.use_gpu)
+    else:
+        ray.init(address=args.address)
+        result = train_tensorflow_mnist(
+            num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
+        )
+    print(result)

python/ray/train/examples/tf/tensorflow_mnist_example.py

@@ -0,0 +1,138 @@
+# This example showcases how to use Tensorflow with Ray Train.
+# Original code:
+# https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import argparse
+import json
+
+import numpy as np
+import tensorflow as tf
+from filelock import FileLock
+
+from ray.air.integrations.keras import ReportCheckpointCallback
+from ray.train import Result, RunConfig, ScalingConfig
+from ray.train.tensorflow import TensorflowTrainer
+
+
+def mnist_dataset(batch_size: int) -> tf.data.Dataset:
+    with FileLock(os.path.expanduser("~/.mnist_lock")):
+        (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
+    # The `x` arrays are in uint8 and have values in the [0, 255] range.
+    # You need to convert them to float32 with values in the [0, 1] range.
+    x_train = x_train / np.float32(255)
+    y_train = y_train.astype(np.int64)
+    train_dataset = (
+        tf.data.Dataset.from_tensor_slices((x_train, y_train))
+        .shuffle(60000)
+        .repeat()
+        .batch(batch_size)
+    )
+    return train_dataset
+
+
+def build_cnn_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=(28, 28)),
+            tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
+            tf.keras.layers.Conv2D(32, 3, activation="relu"),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(10),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+    per_worker_batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+    steps_per_epoch = config.get("steps_per_epoch", 70)
+
+    tf_config = json.loads(os.environ["TF_CONFIG"])
+    num_workers = len(tf_config["cluster"]["worker"])
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    global_batch_size = per_worker_batch_size * num_workers
+    multi_worker_dataset = mnist_dataset(global_batch_size)
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_cnn_model()
+        learning_rate = config.get("lr", 0.001)
+        multi_worker_model.compile(
+            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate),
+            metrics=["accuracy"],
+        )
+
+    history = multi_worker_model.fit(
+        multi_worker_dataset,
+        epochs=epochs,
+        steps_per_epoch=steps_per_epoch,
+        callbacks=[ReportCheckpointCallback()],
+    )
+    results = history.history
+    return results
+
+
+def train_tensorflow_mnist(
+    num_workers: int = 2,
+    use_gpu: bool = False,
+    epochs: int = 4,
+    storage_path: str = None,
+) -> Result:
+    config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+        run_config=RunConfig(storage_path=storage_path),
+    )
+    results = trainer.fit()
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    import ray
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        num_gpus = args.num_workers if args.use_gpu else 0
+        ray.init(num_cpus=4, num_gpus=num_gpus)
+        train_tensorflow_mnist(num_workers=2, use_gpu=args.use_gpu)
+    else:
+        ray.init(address=args.address)
+        train_tensorflow_mnist(
+            num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
+        )

python/ray/train/examples/tf/tensorflow_quick_start.py

@@ -0,0 +1,90 @@
+# ruff: noqa
+# fmt: off
+# isort: skip_file
+
+# __tf_setup_begin__
+import os
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import sys
+import numpy as np
+
+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
+
+def mnist_dataset(batch_size):
+    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
+    # The `x` arrays are in uint8 and have values in the [0, 255] range.
+    # You need to convert them to float32 with values in the [0, 1] range.
+    x_train = x_train / np.float32(255)
+    y_train = y_train.astype(np.int64)
+    train_dataset = tf.data.Dataset.from_tensor_slices(
+        (x_train, y_train)).shuffle(60000).repeat().batch(batch_size)
+    return train_dataset
+
+
+def build_and_compile_cnn_model():
+    model = tf.keras.Sequential([
+        tf.keras.layers.InputLayer(input_shape=(28, 28)),
+        tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
+        tf.keras.layers.Conv2D(32, 3, activation='relu'),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(128, activation='relu'),
+        tf.keras.layers.Dense(10)
+    ])
+    model.compile(
+        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
+        metrics=['accuracy'])
+    return model
+# __tf_setup_end__
+
+# __tf_single_begin__
+def train_func():
+    batch_size = 64
+    single_worker_dataset = mnist_dataset(batch_size)
+    single_worker_model = build_and_compile_cnn_model()
+    single_worker_model.fit(single_worker_dataset, epochs=3, steps_per_epoch=70)
+# __tf_single_end__
+
+# __tf_distributed_begin__
+import json
+import os
+
+def train_func_distributed():
+    per_worker_batch_size = 64
+    # This environment variable will be set by Ray Train.
+    tf_config = json.loads(os.environ['TF_CONFIG'])
+    num_workers = len(tf_config['cluster']['worker'])
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    global_batch_size = per_worker_batch_size * num_workers
+    multi_worker_dataset = mnist_dataset(global_batch_size)
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_and_compile_cnn_model()
+
+    multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70)
+# __tf_distributed_end__
+
+if __name__ == "__main__":
+    # __tf_single_run_begin__
+    train_func()
+    # __tf_single_run_end__
+
+    # __tf_trainer_begin__
+    from ray.train.tensorflow import TensorflowTrainer
+    from ray.train import ScalingConfig
+
+    # For GPU Training, set `use_gpu` to True.
+    use_gpu = False
+
+    trainer = TensorflowTrainer(train_func_distributed, scaling_config=ScalingConfig(num_workers=4, use_gpu=use_gpu))
+
+    trainer.fit()
+    # __tf_trainer_end__

python/ray/train/examples/tf/tensorflow_regression_example.py

@@ -0,0 +1,115 @@
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import argparse
+import sys
+
+import ray
+from ray import train
+from ray.data.preprocessors import Concatenator
+from ray.train import Result, ScalingConfig
+
+if sys.version_info >= (3, 12):
+    # Skip this test in Python 3.12+ because TensorFlow is not supported.
+    sys.exit(0)
+else:
+    import tensorflow as tf
+
+    from ray.train.tensorflow import TensorflowTrainer
+    from ray.train.tensorflow.keras import ReportCheckpointCallback
+
+
+def build_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.InputLayer(input_shape=(100,)),
+            tf.keras.layers.Dense(10),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+    batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+
+    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_absolute_error,
+            metrics=[tf.keras.metrics.mean_squared_error],
+        )
+
+    dataset = train.get_dataset_shard("train")
+
+    results = []
+    for _ in range(epochs):
+        tf_dataset = dataset.to_tf(
+            feature_columns="x", label_columns="y", batch_size=batch_size
+        )
+        history = multi_worker_model.fit(
+            tf_dataset, callbacks=[ReportCheckpointCallback()]
+        )
+        results.append(history.history)
+    return results
+
+
+def train_tensorflow_regression(num_workers: int = 2, use_gpu: bool = False) -> Result:
+    dataset = ray.data.read_csv("s3://anonymous@air-example-data/regression.csv")
+    columns_to_concatenate = [f"x{i:03}" for i in range(100)]
+    preprocessor = Concatenator(columns=columns_to_concatenate, output_column_name="x")
+    dataset = preprocessor.fit_transform(dataset)
+
+    config = {"lr": 1e-3, "batch_size": 32, "epochs": 4}
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=scaling_config,
+        datasets={"train": dataset},
+    )
+    results = trainer.fit()
+    print(results.metrics)
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        num_gpus = args.num_workers if args.use_gpu else 0
+        ray.init(num_cpus=4, num_gpus=num_gpus)
+        result = train_tensorflow_regression(num_workers=2, use_gpu=args.use_gpu)
+    else:
+        ray.init(address=args.address)
+        result = train_tensorflow_regression(
+            num_workers=args.num_workers, use_gpu=args.use_gpu
+        )
+    print(result)

python/ray/train/examples/transformers/transformers_torch_trainer_basic.py

@@ -0,0 +1,78 @@
+import evaluate
+import numpy as np
+
+# Minimal Example adapted from https://huggingface.co/docs/transformers/training
+from datasets import load_dataset
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+)
+
+from ray.train import ScalingConfig
+from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer
+from ray.train.torch import TorchTrainer
+
+
+# [1] Define a training function that includes all your training logic
+# ====================================================================
+def train_func(config):
+    # Datasets
+    dataset = load_dataset("Yelp/yelp_review_full")
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], padding="max_length", truncation=True)
+
+    tokenized_ds = dataset.map(tokenize_function, batched=True)
+
+    small_train_ds = tokenized_ds["train"].shuffle(seed=42).select(range(1000))
+    small_eval_ds = tokenized_ds["test"].shuffle(seed=42).select(range(1000))
+
+    # Model
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", num_labels=5
+    )
+
+    # Evaluation metrics
+    metric = evaluate.load("accuracy")
+
+    def compute_metrics(eval_pred):
+        logits, labels = eval_pred
+        predictions = np.argmax(logits, axis=-1)
+        return metric.compute(predictions=predictions, references=labels)
+
+    # Hugging Face Trainer
+    training_args = TrainingArguments(
+        output_dir="test_trainer", eval_strategy="epoch", report_to="none"
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=small_train_ds,
+        eval_dataset=small_eval_ds,
+        compute_metrics=compute_metrics,
+    )
+
+    # [2] Report metrics and checkpoints to Ray Train
+    # ===============================================
+    trainer.add_callback(RayTrainReportCallback())
+
+    # [3] Prepare your trainer for Ray Data integration
+    # =================================================
+    trainer = prepare_trainer(trainer)
+
+    # Start Training
+    trainer.train()
+
+
+if __name__ == "__main__":
+    # [4] Build a Ray TorchTrainer to launch `train_func` on all workers
+    # ==================================================================
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
+    )
+
+    trainer.fit()

python/ray/train/horovod/__init__.py

@@ -0,0 +1,22 @@
+# isort: off
+try:
+    import horovod  # noqa: F401
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "Horovod isn't installed. To install Horovod with PyTorch support, run 'pip "
+        "install 'horovod[pytorch]''. To install Horovod with TensorFlow support, "
+        "run 'pip install 'horovod[tensorflow]''."
+    )
+# isort: on
+
+from ray.train.horovod.config import HorovodConfig
+from ray.train.horovod.horovod_trainer import HorovodTrainer
+from ray.train.v2._internal.constants import is_v2_enabled
+
+if is_v2_enabled():
+    from ray.train.v2.horovod.horovod_trainer import HorovodTrainer  # noqa: F811
+
+__all__ = ["HorovodConfig", "HorovodTrainer"]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/horovod/config.py

@@ -0,0 +1,163 @@
+import os
+from dataclasses import dataclass
+from typing import Optional, Set
+
+from horovod.ray.runner import Coordinator
+from horovod.ray.utils import detect_nics, nics_to_env_var
+from horovod.runner.common.util import secret, timeout
+
+import ray
+from ray.train._internal.utils import update_env_vars
+from ray.train._internal.worker_group import Worker, WorkerGroup
+from ray.train.backend import Backend, BackendConfig
+from ray.util import PublicAPI
+
+
+@PublicAPI(stability="beta")
+@dataclass
+class HorovodConfig(BackendConfig):
+    """Configurations for Horovod setup.
+
+    See https://github.com/horovod/horovod/blob/master/horovod/runner/common/util/settings.py # noqa: E501
+
+    Args:
+        nics (Optional[Set[str]): Network interfaces that can be used for
+            communication.
+        verbose: Horovod logging verbosity.
+        key (Optional[str]): Secret used for communication between workers.
+        ssh_port (Optional[int]): Port for SSH server running on worker nodes.
+        ssh_identity_file (Optional[str]): Path to the identity file to
+            ssh into different hosts on the cluster.
+        ssh_str (Optional[str]): CAUTION WHEN USING THIS. Private key
+            file contents. Writes the private key to ssh_identity_file.
+        timeout_s: Timeout parameter for Gloo rendezvous.
+        placement_group_timeout_s: Timeout parameter for Ray
+            Placement Group creation. Currently unused.
+    """
+
+    nics: Optional[Set[str]] = None
+    verbose: int = 1
+    key: Optional[str] = None
+    ssh_port: Optional[int] = None
+    ssh_identity_file: Optional[str] = None
+    ssh_str: Optional[str] = None
+    timeout_s: int = 300
+    placement_group_timeout_s: int = 100
+
+    @property
+    def start_timeout(self):
+        return timeout.Timeout(
+            self.timeout_s,
+            message="Timed out waiting for {activity}. Please "
+            "check connectivity between servers. You "
+            "may need to increase the --start-timeout "
+            "parameter if you have too many servers.",
+        )
+
+    def __post_init__(self):
+        if self.ssh_str and not os.path.exists(self.ssh_identity_file):
+            with open(self.ssh_identity_file, "w") as f:
+                os.chmod(self.ssh_identity_file, 0o600)
+                f.write(self.ssh_str)
+
+        if self.key is None:
+            self.key = secret.make_secret_key()
+
+    @property
+    def backend_cls(self):
+        return _HorovodBackend
+
+
+class _HorovodBackend(Backend):
+    share_cuda_visible_devices: bool = True
+
+    def on_start(self, worker_group: WorkerGroup, backend_config: HorovodConfig):
+        # NOTE: Horovod backend uses V1 WorkerGroup directly instead of BaseWorkerGroup
+        # because it requires direct access to worker metadata (node_id, hostname) that is
+        # specific to the V1 implementation. Horovod does not support V2 WorkerGroup.
+
+        # TODO(matt): Implement placement group strategies in BackendExecutor.
+
+        # Initialize workers with Horovod environment variables
+        setup_futures = []
+        for rank in range(len(worker_group)):
+            worker_node_id = worker_group.workers[rank].metadata.node_id
+            setup_futures.append(
+                worker_group.execute_single_async(
+                    rank,
+                    _init_env_vars,
+                    rank,
+                    len(worker_group),
+                    worker_node_id,
+                )
+            )
+        ray.get(setup_futures)
+
+        # Use Horovod Ray Coordinator
+        # backend_config as settings
+        self.coordinator = Coordinator(backend_config)
+
+        # Get all the hostnames of all workers
+        node_ids = [w.metadata.node_id for w in worker_group.workers]
+        hostnames = [w.metadata.hostname for w in worker_group.workers]
+        # Register each hostname to the coordinator. assumes the hostname
+        # ordering is the same.
+        for rank, (hostname, node_id) in enumerate(zip(hostnames, node_ids)):
+            self.coordinator.register(hostname, node_id, rank)
+        all_info = self.coordinator.finalize_registration()
+
+        setup_futures = []
+        for rank, local_cross_env_var in all_info.items():
+            setup_futures.append(
+                worker_group.execute_single_async(
+                    rank, update_env_vars, local_cross_env_var
+                )
+            )
+        ray.get(setup_futures)
+
+        coordinator_envs = self.coordinator.establish_rendezvous()
+
+        # Get one worker from each host/node.
+        node_worker_indexes = [node_ids.index(node_id) for node_id in set(node_ids)]
+        node_workers = [
+            _HorovodWorkerWrapper(worker_group.workers[worker_index])
+            for worker_index in node_worker_indexes
+        ]
+        assert len(node_workers) == len(self.coordinator.hostnames)
+
+        nics = detect_nics(
+            backend_config,
+            all_host_names=list(self.coordinator.hostnames),
+            node_workers=node_workers,
+        )
+        coordinator_envs.update(nics_to_env_var(nics))
+
+        worker_group.execute(update_env_vars, coordinator_envs)
+
+
+def _init_env_vars(world_rank: int, world_size: int, node_id: str):
+    """Initialize Horovod environment variables."""
+    os.environ["HOROVOD_HOSTNAME"] = node_id
+    os.environ["HOROVOD_RANK"] = str(world_rank)
+    os.environ["HOROVOD_SIZE"] = str(world_size)
+
+
+# TODO(tgaddair): temporary workaround for Horovod's worker discovery logic,
+#  which requires passing in an extra parameter as part of the RayExecutor
+#  API. This will be removed in the future as we migrate more of the
+#  RayExecutor utils into Ray Train.
+#  See: https://github.com/horovod/horovod/blob/v0.23.0/horovod/ray/driver_service.py#L9 # noqa: E501
+@dataclass
+class _HorovodWorkerWrapper:
+    w: Worker
+
+    @property
+    def execute(self):
+        w = self.w
+
+        class ExecuteHandle:
+            def remote(self, func, *args, **kwargs):
+                _ = None
+                return w.actor._RayTrainWorker__execute.remote(func, _, *args, **kwargs)
+
+        return ExecuteHandle()

python/ray/train/horovod/horovod_trainer.py

@@ -0,0 +1,197 @@
+from typing import Any, Callable, Dict, Optional, Union
+
+from ray.air.config import RunConfig, ScalingConfig
+from ray.train import Checkpoint, DataConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.horovod.config import HorovodConfig
+from ray.train.trainer import GenDataset
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="beta")
+class HorovodTrainer(DataParallelTrainer):
+    """A Trainer for data parallel Horovod training.
+
+    This Trainer runs the function ``train_loop_per_worker`` on multiple Ray
+    Actors. These actors already have the necessary Horovod setup already
+    configured for distributed Horovod training.
+
+    The ``train_loop_per_worker`` function is expected to take in either 0 or 1
+    arguments:
+
+    .. testcode::
+
+        def train_loop_per_worker():
+            ...
+
+    .. testcode::
+
+        def train_loop_per_worker(config: Dict):
+            ...
+
+    If ``train_loop_per_worker`` accepts an argument, then
+    ``train_loop_config`` will be passed in as the argument. This is useful if you
+    want to tune the values in ``train_loop_config`` as hyperparameters.
+
+    If the ``datasets`` dict contains a training dataset (denoted by
+    the "train" key), then it will be split into multiple dataset
+    shards that can then be accessed by ``ray.train.get_dataset_shard("train")`` inside
+    ``train_loop_per_worker``. All the other datasets will not be split and
+    ``ray.train.get_dataset_shard(...)`` will return the entire Dataset.
+
+    Inside the ``train_loop_per_worker`` function, you can use any of the
+    :ref:`Ray Train loop methods <train-loop-api>`.
+
+    .. testcode::
+
+        from ray import train
+
+        def train_loop_per_worker():
+            # Report intermediate results for callbacks or logging and
+            # checkpoint data.
+            train.report(...)
+
+            # Returns dict of last saved checkpoint.
+            train.get_checkpoint()
+
+            # Returns the Dataset shard for the given key.
+            train.get_dataset_shard("my_dataset")
+
+            # Returns the total number of workers executing training.
+            train.get_context().get_world_size()
+
+            # Returns the rank of this worker.
+            train.get_context().get_world_rank()
+
+            # Returns the rank of the worker on the current node.
+            train.get_context().get_local_rank()
+
+    Any returns from the ``train_loop_per_worker`` will be discarded and not
+    used or persisted anywhere.
+
+    Example:
+
+
+    .. testcode::
+        :skipif: True
+
+        import os
+        import tempfile
+
+        import ray
+        import horovod.torch as hvd
+        import torch
+        import torch.nn as nn
+
+        from ray import train
+        import ray.train.torch  # Need this to use `train.torch.get_device()`
+        from ray.train import Checkpoint, ScalingConfig
+        from ray.train.horovod import HorovodTrainer
+
+        # If using GPUs, set this to True.
+        use_gpu = False
+
+        input_size = 1
+        layer_size = 15
+        output_size = 1
+        num_epochs = 3
+
+        class NeuralNetwork(nn.Module):
+            def __init__(self):
+                super(NeuralNetwork, self).__init__()
+                self.layer1 = nn.Linear(input_size, layer_size)
+                self.relu = nn.ReLU()
+                self.layer2 = nn.Linear(layer_size, output_size)
+            def forward(self, input):
+                return self.layer2(self.relu(self.layer1(input)))
+
+        def train_loop_per_worker():
+            hvd.init()
+            dataset_shard = train.get_dataset_shard("train")
+            model = NeuralNetwork()
+            device = train.torch.get_device()
+            model.to(device)
+            loss_fn = nn.MSELoss()
+            lr_scaler = 1
+            optimizer = torch.optim.SGD(model.parameters(), lr=0.1 * lr_scaler)
+            # Horovod: wrap optimizer with DistributedOptimizer.
+            optimizer = hvd.DistributedOptimizer(
+                optimizer,
+                named_parameters=model.named_parameters(),
+                op=hvd.Average,
+            )
+            for epoch in range(num_epochs):
+                model.train()
+                for batch in dataset_shard.iter_torch_batches(
+                    batch_size=32, dtypes=torch.float
+                ):
+                    inputs, labels = torch.unsqueeze(batch["x"], 1), batch["y"]
+                    outputs = model(inputs)
+                    loss = loss_fn(outputs, labels)
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
+                    print(f"epoch: {epoch}, loss: {loss.item()}")
+
+                # Save a model checkpoint at the end of each epoch
+                with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                    ckpt_path = os.path.join(temp_checkpoint_dir, "model.pt")
+                    torch.save(model.state_dict(), ckpt_path)
+                    train.report(
+                        {"loss": loss.item(), "epoch": epoch},
+                        checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
+                    )
+
+        train_dataset = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
+        scaling_config = ScalingConfig(num_workers=3, use_gpu=use_gpu)
+        trainer = HorovodTrainer(
+            train_loop_per_worker=train_loop_per_worker,
+            scaling_config=scaling_config,
+            datasets={"train": train_dataset},
+        )
+        result = trainer.fit()
+
+    Args:
+        train_loop_per_worker: The training function to execute.
+            This can either take in no arguments or a ``config`` dict.
+        train_loop_config: Configurations to pass into
+            ``train_loop_per_worker`` if it accepts an argument.
+        horovod_config: Configuration for setting up the Horovod backend.
+            If set to None, use the default configuration. This replaces the
+            ``backend_config`` arg of ``DataParallelTrainer``.
+        scaling_config: Configuration for how to scale data parallel training.
+        dataset_config: Configuration for dataset ingest.
+        run_config: Configuration for the execution of the training run.
+        datasets: Any Datasets to use for training. Use
+            the key "train" to denote which dataset is the training
+            dataset.
+        metadata: Dict that should be made available via
+            `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
+            for checkpoints saved from this Trainer. Must be JSON-serializable.
+        resume_from_checkpoint: A checkpoint to resume training from.
+    """
+
+    def __init__(
+        self,
+        train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]],
+        *,
+        train_loop_config: Optional[Dict] = None,
+        horovod_config: Optional[HorovodConfig] = None,
+        scaling_config: Optional[ScalingConfig] = None,
+        dataset_config: Optional[DataConfig] = None,
+        run_config: Optional[RunConfig] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        resume_from_checkpoint: Optional[Checkpoint] = None,
+    ):
+        super().__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            train_loop_config=train_loop_config,
+            backend_config=horovod_config or HorovodConfig(),
+            scaling_config=scaling_config,
+            dataset_config=dataset_config,
+            run_config=run_config,
+            datasets=datasets,
+            resume_from_checkpoint=resume_from_checkpoint,
+            metadata=metadata,
+        )

python/ray/train/huggingface/transformers/__init__.py

@@ -0,0 +1,12 @@
+from ray.train.huggingface.transformers._transformers_utils import (
+    RayTrainReportCallback,
+    prepare_trainer,
+)
+
+__all__ = [
+    "RayTrainReportCallback",
+    "prepare_trainer",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/huggingface/transformers/_transformers_utils.py

@@ -0,0 +1,150 @@
+import logging
+import shutil
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Iterator, Optional, Type, Union
+
+from torch.utils.data import DataLoader, Dataset, IterableDataset
+
+import ray
+from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
+from ray.train import Checkpoint
+from ray.util import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+TRANSFORMERS_IMPORT_ERROR: Optional[ImportError] = None
+
+try:
+    import transformers.trainer
+    from transformers import Trainer
+    from transformers.trainer_callback import TrainerCallback
+except ImportError as e:
+    TRANSFORMERS_IMPORT_ERROR = e
+    TrainerCallback = object
+
+
+@PublicAPI(stability="beta")
+class RayTrainReportCallback(TrainerCallback):
+    """A simple callback to report checkpoints and metrics to Ray Train.
+
+    This callback is a subclass of `transformers.TrainerCallback
+    <https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback>`_
+    and overrides the `TrainerCallback.on_save()` method. After
+    a new checkpoint get saved, it fetches the latest metric dictionary
+    from `TrainerState.log_history` and reports it with the latest checkpoint
+    to Ray Train.
+
+    Checkpoints will be saved in the following structure::
+
+        checkpoint_00000*/   Ray Train Checkpoint
+        └─ checkpoint/       Hugging Face Transformers Checkpoint
+
+    For customized reporting and checkpointing logic, implement your own
+    `transformers.TrainerCallback` following this user
+    guide: :ref:`Saving and Loading Checkpoints <train-dl-saving-checkpoints>`.
+
+    Note that users should ensure that the logging, evaluation, and saving frequencies
+    are properly configured so that the monitoring metric is always up-to-date
+    when `transformers.Trainer` saves a checkpoint.
+
+    Suppose the monitoring metric is reported from evaluation stage:
+
+    Some valid configurations:
+        - evaluation_strategy == save_strategy == "epoch"
+        - evaluation_strategy == save_strategy == "steps", save_steps % eval_steps == 0
+
+    Some invalid configurations:
+        - evaluation_strategy != save_strategy
+        - evaluation_strategy == save_strategy == "steps", save_steps % eval_steps != 0
+
+    """
+
+    CHECKPOINT_NAME = "checkpoint"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        record_extra_usage_tag(TagKey.TRAIN_TRANSFORMERS_RAYTRAINREPORTCALLBACK, "1")
+
+    def on_save(self, args, state, control, **kwargs):
+        """Event called after a checkpoint save."""
+        with TemporaryDirectory() as tmpdir:
+            # Aggregate all the logged metrics
+            metrics = {}
+            for log in state.log_history:
+                metrics.update(log)
+
+            # Copy ckpt files and construct a Ray Train Checkpoint
+            source_ckpt_path = transformers.trainer.get_last_checkpoint(args.output_dir)
+            if source_ckpt_path is not None:
+                target_ckpt_path = Path(tmpdir, self.CHECKPOINT_NAME).as_posix()
+                shutil.copytree(source_ckpt_path, target_ckpt_path)
+                checkpoint = Checkpoint.from_directory(tmpdir)
+            else:
+                checkpoint = None
+
+            # Report latest metrics and checkpoint to Ray Train
+            ray.train.report(metrics=metrics, checkpoint=checkpoint)
+
+
+class RayTorchIterableDataset(IterableDataset):
+    """Wrapper class for ray data iterables."""
+
+    def __init__(self, data_iterable) -> None:
+        super().__init__()
+        self.data_iterable = data_iterable
+
+    def __iter__(self) -> Iterator:
+        return iter(self.data_iterable)
+
+
+@PublicAPI(stability="beta")
+def prepare_trainer(trainer: "Trainer") -> "Trainer":
+    """Prepare your HuggingFace Transformer Trainer for Ray Train.
+
+    This utility function enable the trainer integrates with Ray Data Integration.
+    Internally, it overrides the `get_train_dataloader` and `get_eval_dataloader`
+    methods and inject the data integration logics if the `train_dataset` and
+    `eval_dataset` are Ray Data Iterables.
+    """
+    from ray.data.iterator import _IterableFromIterator
+
+    if TRANSFORMERS_IMPORT_ERROR is not None:
+        raise TRANSFORMERS_IMPORT_ERROR
+
+    base_trainer_class: Type[transformers.trainer.Trainer] = trainer.__class__
+
+    class RayTransformersTrainer(base_trainer_class):
+        """A Wrapper of `transformers.Trainer` for Ray Data Integration."""
+
+        def get_train_dataloader(self) -> DataLoader:
+            if isinstance(self.train_dataset, _IterableFromIterator):
+                dataset = RayTorchIterableDataset(self.train_dataset)
+                return DataLoader(dataset, batch_size=1, collate_fn=lambda x: x[0])
+            else:
+                return super().get_train_dataloader()
+
+        def get_eval_dataloader(
+            self, eval_dataset: Optional[Union[str, Dataset]] = None
+        ) -> DataLoader:
+            if eval_dataset is None:
+                eval_dataset = self.eval_dataset
+
+            if (
+                isinstance(eval_dataset, str)
+                and isinstance(self.eval_dataset, dict)
+                and isinstance(self.eval_dataset[eval_dataset], _IterableFromIterator)
+            ):
+                dataset = RayTorchIterableDataset(self.eval_dataset[eval_dataset])
+                return DataLoader(dataset, batch_size=1, collate_fn=lambda x: x[0])
+            elif isinstance(eval_dataset, _IterableFromIterator):
+                dataset = RayTorchIterableDataset(eval_dataset)
+                return DataLoader(dataset, batch_size=1, collate_fn=lambda x: x[0])
+            else:
+                return super().get_eval_dataloader(eval_dataset)
+
+    trainer.__class__ = RayTransformersTrainer
+
+    record_extra_usage_tag(TagKey.TRAIN_TRANSFORMERS_PREPARE_TRAINER, "1")
+    return trainer

python/ray/train/lightgbm/__init__.py

@@ -0,0 +1,23 @@
+from ray.train.lightgbm._lightgbm_utils import (
+    RayTrainReportCallback,
+    normalize_pandas_for_lightgbm,
+)
+from ray.train.lightgbm.config import LightGBMConfig, get_network_params
+from ray.train.lightgbm.lightgbm_checkpoint import LightGBMCheckpoint
+from ray.train.lightgbm.lightgbm_trainer import LightGBMTrainer
+from ray.train.v2._internal.constants import is_v2_enabled
+
+if is_v2_enabled():
+    from ray.train.v2.lightgbm.lightgbm_trainer import LightGBMTrainer  # noqa: F811
+
+__all__ = [
+    "RayTrainReportCallback",
+    "LightGBMCheckpoint",
+    "LightGBMTrainer",
+    "LightGBMConfig",
+    "get_network_params",
+    "normalize_pandas_for_lightgbm",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/lightgbm/_lightgbm_utils.py

@@ -0,0 +1,256 @@
+import tempfile
+from abc import abstractmethod
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
+
+from lightgbm.basic import Booster
+from lightgbm.callback import CallbackEnv
+
+import ray.train
+from ray.train import Checkpoint
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+@PublicAPI(stability="alpha")
+def normalize_pandas_for_lightgbm(df: "pd.DataFrame") -> "pd.DataFrame":
+    """Map Arrow-backed pandas dtypes to NumPy-nullable equivalents.
+
+    LightGBM's pandas input validation rejects Arrow-backed dtypes like
+    ``int64[pyarrow]``. Since Ray Data 2.56, ``Dataset.to_pandas()`` preserves
+    Arrow-backed dtypes when the source was Arrow, so callers passing the
+    resulting frame to ``lightgbm.Dataset`` must normalize first.
+
+    This helper is a faster alternative to
+    ``df.convert_dtypes(dtype_backend="numpy_nullable")``:
+
+    - It maps dtypes mechanically rather than scanning every value.
+    - It only touches ``pd.ArrowDtype`` columns. NumPy-backed columns (e.g.
+      from ``ray.data.from_pandas`` shards) keep their original buffers.
+
+    Only numeric and boolean Arrow dtypes are remapped. Other Arrow dtypes
+    (string, decimal, timestamp) are left as-is; LightGBM doesn't accept them
+    as features anyway.
+
+    Args:
+        df: The pandas DataFrame to normalize.
+
+    Returns:
+        A DataFrame with Arrow-backed numeric/boolean columns replaced by
+        NumPy-nullable equivalents. Other columns are returned unchanged.
+    """
+    import pandas as pd
+    import pyarrow as pa
+
+    dtype_mapping = {}
+    for column, dtype in df.dtypes.items():
+        if not isinstance(dtype, pd.ArrowDtype):
+            continue
+        arrow_dtype = dtype.pyarrow_dtype
+        if pa.types.is_signed_integer(arrow_dtype):
+            dtype_mapping[column] = f"Int{arrow_dtype.bit_width}"
+        elif pa.types.is_unsigned_integer(arrow_dtype):
+            dtype_mapping[column] = f"UInt{arrow_dtype.bit_width}"
+        elif pa.types.is_floating(arrow_dtype):
+            dtype_mapping[column] = f"Float{arrow_dtype.bit_width}"
+        elif pa.types.is_boolean(arrow_dtype):
+            dtype_mapping[column] = "boolean"
+    if dtype_mapping:
+        df = df.astype(dtype_mapping, copy=False)
+    return df
+
+
+class RayReportCallback:
+    CHECKPOINT_NAME = "model.txt"
+
+    def __init__(
+        self,
+        metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
+        filename: str = CHECKPOINT_NAME,
+        frequency: int = 0,
+        checkpoint_at_end: bool = True,
+        results_postprocessing_fn: Optional[
+            Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]]
+        ] = None,
+    ):
+        if isinstance(metrics, str):
+            metrics = [metrics]
+        self._metrics = metrics
+        self._filename = filename
+        self._frequency = frequency
+        self._checkpoint_at_end = checkpoint_at_end
+        self._results_postprocessing_fn = results_postprocessing_fn
+
+    @classmethod
+    def get_model(
+        cls, checkpoint: Checkpoint, filename: str = CHECKPOINT_NAME
+    ) -> Booster:
+        """Retrieve the model stored in a checkpoint reported by this callback.
+
+        Args:
+            checkpoint: The checkpoint object returned by a training run.
+                The checkpoint should be saved by an instance of this callback.
+            filename: The filename to load the model from, which should match
+                the filename used when creating the callback.
+        Returns:
+            The model loaded from the checkpoint.
+        """
+        with checkpoint.as_directory() as checkpoint_path:
+            return Booster(model_file=Path(checkpoint_path, filename).as_posix())
+
+    def _get_report_dict(self, evals_log: Dict[str, Dict[str, list]]) -> dict:
+        result_dict = flatten_dict(evals_log, delimiter="-")
+        if not self._metrics:
+            report_dict = result_dict
+        else:
+            report_dict = {}
+            for key in self._metrics:
+                if isinstance(self._metrics, dict):
+                    metric = self._metrics[key]
+                else:
+                    metric = key
+                report_dict[key] = result_dict[metric]
+        if self._results_postprocessing_fn:
+            report_dict = self._results_postprocessing_fn(report_dict)
+        return report_dict
+
+    def _get_eval_result(self, env: CallbackEnv) -> dict:
+        eval_result = {}
+        for entry in env.evaluation_result_list:
+            data_name, eval_name, result = entry[0:3]
+            if len(entry) > 4:
+                stdv = entry[4]
+                suffix = "-mean"
+            else:
+                stdv = None
+                suffix = ""
+            if data_name not in eval_result:
+                eval_result[data_name] = {}
+            eval_result[data_name][eval_name + suffix] = result
+            if stdv is not None:
+                eval_result[data_name][eval_name + "-stdv"] = stdv
+        return eval_result
+
+    @abstractmethod
+    def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
+        """Get checkpoint from model.
+
+        This method needs to be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
+        """Save checkpoint and report metrics corresonding to this checkpoint.
+
+        This method needs to be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _report_metrics(self, report_dict: Dict):
+        """Report Metrics.
+
+        This method needs to be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    def __call__(self, env: CallbackEnv) -> None:
+        eval_result = self._get_eval_result(env)
+        report_dict = self._get_report_dict(eval_result)
+
+        # Ex: if frequency=2, checkpoint_at_end=True and num_boost_rounds=11,
+        # you will checkpoint at iterations 1, 3, 5, ..., 9, and 10 (checkpoint_at_end)
+        # (iterations count from 0)
+        on_last_iter = env.iteration == env.end_iteration - 1
+        should_checkpoint_at_end = on_last_iter and self._checkpoint_at_end
+        should_checkpoint_with_frequency = (
+            self._frequency != 0 and (env.iteration + 1) % self._frequency == 0
+        )
+        should_checkpoint = should_checkpoint_at_end or should_checkpoint_with_frequency
+
+        if should_checkpoint:
+            self._save_and_report_checkpoint(report_dict, env.model)
+        else:
+            self._report_metrics(report_dict)
+
+
+@PublicAPI(stability="beta")
+class RayTrainReportCallback(RayReportCallback):
+    """Creates a callback that reports metrics and checkpoints model.
+
+    Args:
+        metrics: Metrics to report. If this is a list,
+            each item should be a metric key reported by LightGBM,
+            and it will be reported to Ray Train/Tune under the same name.
+            This can also be a dict of {<key-to-report>: <lightgbm-metric-key>},
+            which can be used to rename LightGBM default metrics.
+        filename: Customize the saved checkpoint file type by passing
+            a filename. Defaults to "model.txt".
+        frequency: How often to save checkpoints, in terms of iterations.
+            Defaults to 0 (no checkpoints are saved during training).
+        checkpoint_at_end: Whether or not to save a checkpoint at the end of training.
+        results_postprocessing_fn: An optional Callable that takes in
+            the metrics dict that will be reported (after it has been flattened)
+            and returns a modified dict.
+
+    Examples
+    --------
+
+    Reporting checkpoints and metrics to Ray Tune when running many
+    independent LightGBM trials (without data parallelism within a trial).
+
+    .. testcode::
+        :skipif: True
+
+        import lightgbm
+
+        from ray.train.lightgbm import RayTrainReportCallback
+
+        config = {
+            # ...
+            "metric": ["binary_logloss", "binary_error"],
+        }
+
+        # Report only log loss to Tune after each validation epoch.
+        bst = lightgbm.train(
+            ...,
+            callbacks=[
+                RayTrainReportCallback(
+                    metrics={"loss": "eval-binary_logloss"}, frequency=1
+                )
+            ],
+        )
+
+    Loading a model from a checkpoint reported by this callback.
+
+    .. testcode::
+        :skipif: True
+
+        from ray.train.lightgbm import RayTrainReportCallback
+
+        # Get a `Checkpoint` object that is saved by the callback during training.
+        result = trainer.fit()
+        booster = RayTrainReportCallback.get_model(result.checkpoint)
+
+    """
+
+    @contextmanager
+    def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]:
+        if ray.train.get_context().get_world_rank() in (0, None):
+            with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                model.save_model(Path(temp_checkpoint_dir, self._filename).as_posix())
+                yield Checkpoint.from_directory(temp_checkpoint_dir)
+        else:
+            yield None
+
+    def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster):
+        with self._get_checkpoint(model=model) as checkpoint:
+            ray.train.report(report_dict, checkpoint=checkpoint)
+
+    def _report_metrics(self, report_dict: Dict):
+        ray.train.report(report_dict)

python/ray/train/lightgbm/config.py

@@ -0,0 +1,98 @@
+import logging
+import threading
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import ray
+from ray._common.network_utils import build_address
+from ray.train._internal.base_worker_group import BaseWorkerGroup
+from ray.train._internal.utils import get_address_and_port
+from ray.train.backend import Backend, BackendConfig
+from ray.train.v2._internal.util import TrainingFramework
+
+logger = logging.getLogger(__name__)
+
+
+# Global LightGBM distributed network configuration for each worker process.
+_lightgbm_network_params: Optional[Dict[str, Any]] = None
+_lightgbm_network_params_lock = threading.Lock()
+
+
+def get_network_params() -> Dict[str, Any]:
+    """Returns the network parameters to enable LightGBM distributed training."""
+    global _lightgbm_network_params
+
+    with _lightgbm_network_params_lock:
+        if not _lightgbm_network_params:
+            logger.warning(
+                "`ray.train.lightgbm.get_network_params` was called outside "
+                "the context of a `ray.train.lightgbm.LightGBMTrainer`. "
+                "The current process has no knowledge of the distributed training "
+                "worker group, so this method will return an empty dict. "
+                "Please call this within the training loop of a "
+                "`ray.train.lightgbm.LightGBMTrainer`. "
+                "If you are in fact calling this within a `LightGBMTrainer`, "
+                "this is unexpected: please file a bug report to the Ray Team."
+            )
+            return {}
+
+        return _lightgbm_network_params.copy()
+
+
+def _set_network_params(
+    num_machines: int,
+    local_listen_port: int,
+    machines: str,
+):
+    global _lightgbm_network_params
+
+    with _lightgbm_network_params_lock:
+        assert (
+            _lightgbm_network_params is None
+        ), "LightGBM network params are already initialized."
+        _lightgbm_network_params = dict(
+            num_machines=num_machines,
+            local_listen_port=local_listen_port,
+            machines=machines,
+        )
+
+
+@dataclass
+class LightGBMConfig(BackendConfig):
+    """Configuration for LightGBM distributed data-parallel training setup.
+
+    See the LightGBM docs for more information on the "network parameters"
+    that Ray Train sets up for you:
+    https://lightgbm.readthedocs.io/en/latest/Parameters.html#network-parameters
+    """
+
+    @property
+    def backend_cls(self):
+        return _LightGBMBackend
+
+    @property
+    def framework(self):
+        return TrainingFramework.LIGHTGBM
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {}
+
+
+class _LightGBMBackend(Backend):
+    def on_training_start(
+        self, worker_group: BaseWorkerGroup, backend_config: LightGBMConfig
+    ):
+        node_ips_and_ports = worker_group.execute(get_address_and_port)
+        ports = [port for _, port in node_ips_and_ports]
+        machines = ",".join(
+            [build_address(node_ip, port) for node_ip, port in node_ips_and_ports]
+        )
+        num_machines = len(worker_group)
+        ray.get(
+            [
+                worker_group.execute_single_async(
+                    rank, _set_network_params, num_machines, ports[rank], machines
+                )
+                for rank in range(len(worker_group))
+            ]
+        )

python/ray/train/lightgbm/lightgbm_checkpoint.py

@@ -0,0 +1,72 @@
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+import lightgbm
+
+from ray.train._internal.framework_checkpoint import FrameworkCheckpoint
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.preprocessor import Preprocessor
+
+
+@PublicAPI(stability="beta")
+class LightGBMCheckpoint(FrameworkCheckpoint):
+    """A :py:class:`~ray.train.Checkpoint` with LightGBM-specific functionality."""
+
+    MODEL_FILENAME = "model.txt"
+
+    @classmethod
+    def from_model(
+        cls,
+        booster: lightgbm.Booster,
+        *,
+        preprocessor: Optional["Preprocessor"] = None,
+        path: Optional[str] = None,
+    ) -> "LightGBMCheckpoint":
+        """Create a :py:class:`~ray.train.Checkpoint` that stores a LightGBM model.
+
+        Args:
+            booster: The LightGBM model to store in the checkpoint.
+            preprocessor: A fitted preprocessor to be applied before inference.
+            path: The path to the directory where the checkpoint file will be saved.
+                This should start as an empty directory, since the *entire*
+                directory will be treated as the checkpoint when reported.
+                By default, a temporary directory will be created.
+
+        Returns:
+            An :py:class:`LightGBMCheckpoint` containing the specified ``Estimator``.
+
+        Examples:
+            .. testcode::
+
+                import lightgbm
+                import numpy as np
+                from ray.train.lightgbm import LightGBMCheckpoint
+
+                train_X = np.array([[1, 2], [3, 4]])
+                train_y = np.array([0, 1])
+
+                model = lightgbm.LGBMClassifier().fit(train_X, train_y)
+                checkpoint = LightGBMCheckpoint.from_model(model.booster_)
+        """
+        checkpoint_path = Path(path or tempfile.mkdtemp())
+
+        if not checkpoint_path.is_dir():
+            raise ValueError(f"`path` must be a directory, but got: {checkpoint_path}")
+
+        booster.save_model(checkpoint_path.joinpath(cls.MODEL_FILENAME).as_posix())
+
+        checkpoint = cls.from_directory(checkpoint_path.as_posix())
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+
+        return checkpoint
+
+    def get_model(self) -> lightgbm.Booster:
+        """Retrieve the LightGBM model stored in this checkpoint."""
+        with self.as_directory() as checkpoint_path:
+            return lightgbm.Booster(
+                model_file=Path(checkpoint_path, self.MODEL_FILENAME).as_posix()
+            )

python/ray/train/lightgbm/lightgbm_trainer.py

@@ -0,0 +1,326 @@
+import logging
+from functools import partial
+from typing import Any, Callable, Dict, Optional, Union
+
+import lightgbm
+
+import ray
+from ray.train import Checkpoint
+from ray.train.constants import TRAIN_DATASET_KEY
+from ray.train.lightgbm._lightgbm_utils import (
+    RayTrainReportCallback,
+    normalize_pandas_for_lightgbm,
+)
+from ray.train.lightgbm.config import LightGBMConfig
+from ray.train.lightgbm.v2 import LightGBMTrainer as SimpleLightGBMTrainer
+from ray.train.trainer import GenDataset
+from ray.train.utils import _log_deprecation_warning
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+LEGACY_LIGHTGBM_TRAINER_DEPRECATION_MESSAGE = (
+    "Passing in `lightgbm.train` kwargs such as `params`, `num_boost_round`, "
+    "`label_column`, etc. to `LightGBMTrainer` is deprecated "
+    "in favor of the new API which accepts a `train_loop_per_worker` argument, "
+    "similar to the other DataParallelTrainer APIs (ex: TorchTrainer). "
+    "See this issue for more context: "
+    "https://github.com/ray-project/ray/issues/50042"
+)
+
+
+def _lightgbm_train_fn_per_worker(
+    config: dict,
+    label_column: str,
+    num_boost_round: int,
+    dataset_keys: set,
+    lightgbm_train_kwargs: dict,
+):
+    checkpoint = ray.train.get_checkpoint()
+    starting_model = None
+    remaining_iters = num_boost_round
+    if checkpoint:
+        starting_model = RayTrainReportCallback.get_model(checkpoint)
+        starting_iter = starting_model.current_iteration()
+        remaining_iters = num_boost_round - starting_iter
+        logger.info(
+            f"Model loaded from checkpoint will train for "
+            f"additional {remaining_iters} iterations (trees) in order "
+            "to achieve the target number of iterations "
+            f"({num_boost_round=})."
+        )
+
+    train_ds_iter = ray.train.get_dataset_shard(TRAIN_DATASET_KEY)
+    train_df = normalize_pandas_for_lightgbm(train_ds_iter.materialize().to_pandas())
+
+    eval_ds_iters = {
+        k: ray.train.get_dataset_shard(k)
+        for k in dataset_keys
+        if k != TRAIN_DATASET_KEY
+    }
+    eval_dfs = {
+        k: normalize_pandas_for_lightgbm(d.materialize().to_pandas())
+        for k, d in eval_ds_iters.items()
+    }
+
+    train_X, train_y = train_df.drop(label_column, axis=1), train_df[label_column]
+    train_set = lightgbm.Dataset(train_X, label=train_y)
+
+    # NOTE: Include the training dataset in the evaluation datasets.
+    # This allows `train-*` metrics to be calculated and reported.
+    valid_sets = [train_set]
+    valid_names = [TRAIN_DATASET_KEY]
+
+    for eval_name, eval_df in eval_dfs.items():
+        eval_X, eval_y = eval_df.drop(label_column, axis=1), eval_df[label_column]
+        valid_sets.append(lightgbm.Dataset(eval_X, label=eval_y))
+        valid_names.append(eval_name)
+
+    # Add network params of the worker group to enable distributed training.
+    config.update(ray.train.lightgbm.get_network_params())
+    config.setdefault("tree_learner", "data_parallel")
+    config.setdefault("pre_partition", True)
+
+    lightgbm.train(
+        params=config,
+        train_set=train_set,
+        num_boost_round=remaining_iters,
+        valid_sets=valid_sets,
+        valid_names=valid_names,
+        init_model=starting_model,
+        **lightgbm_train_kwargs,
+    )
+
+
+@PublicAPI(stability="beta")
+class LightGBMTrainer(SimpleLightGBMTrainer):
+    """A Trainer for distributed data-parallel LightGBM training.
+
+    Example:
+
+        .. testcode::
+            :skipif: True
+
+            import lightgbm
+
+            import ray.data
+            import ray.train
+            from ray.train.lightgbm import (
+                LightGBMTrainer,
+                RayTrainReportCallback,
+                normalize_pandas_for_lightgbm,
+            )
+
+            def train_fn_per_worker(config: dict):
+                # (Optional) Add logic to resume training state from a checkpoint.
+                # ray.train.get_checkpoint()
+
+                # 1. Get the dataset shard for the worker and convert to a `lightgbm.Dataset`
+                train_ds_iter, eval_ds_iter = (
+                    ray.train.get_dataset_shard("train"),
+                    ray.train.get_dataset_shard("validation"),
+                )
+                train_ds, eval_ds = train_ds_iter.materialize(), eval_ds_iter.materialize()
+                train_df = normalize_pandas_for_lightgbm(train_ds.to_pandas())
+                eval_df = normalize_pandas_for_lightgbm(eval_ds.to_pandas())
+                train_X, train_y = train_df.drop("y", axis=1), train_df["y"]
+                eval_X, eval_y = eval_df.drop("y", axis=1), eval_df["y"]
+                dtrain = lightgbm.Dataset(train_X, label=train_y)
+                deval = lightgbm.Dataset(eval_X, label=eval_y)
+
+                params = {
+                    "objective": "regression",
+                    "metric": "l2",
+                    "learning_rate": 1e-4,
+                    "subsample": 0.5,
+                    "max_depth": 2,
+                    # Adding the line below is the only change needed
+                    # for your `lgb.train` call!
+                    **ray.train.lightgbm.get_network_params(),
+                }
+
+                # 2. Do distributed data-parallel training.
+                # Ray Train sets up the necessary coordinator processes and
+                # environment variables for your workers to communicate with each other.
+                bst = lightgbm.train(
+                    params,
+                    train_set=dtrain,
+                    valid_sets=[deval],
+                    valid_names=["validation"],
+                    num_boost_round=10,
+                    callbacks=[RayTrainReportCallback()],
+                )
+
+            train_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
+            eval_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(16)])
+            trainer = LightGBMTrainer(
+                train_fn_per_worker,
+                datasets={"train": train_ds, "validation": eval_ds},
+                scaling_config=ray.train.ScalingConfig(num_workers=4),
+            )
+            result = trainer.fit()
+            booster = RayTrainReportCallback.get_model(result.checkpoint)
+
+    Args:
+        train_loop_per_worker: The training function to execute on each worker.
+            This function can either take in zero arguments or a single ``Dict``
+            argument which is set by defining ``train_loop_config``.
+            Within this function you can use any of the
+            :ref:`Ray Train Loop utilities <train-loop-api>`.
+        train_loop_config: A configuration ``Dict`` to pass in as an argument to
+            ``train_loop_per_worker``.
+            This is typically used for specifying hyperparameters.
+        lightgbm_config: The configuration for setting up the distributed lightgbm
+            backend. Defaults to using the "rabit" backend.
+            See :class:`~ray.train.lightgbm.LightGBMConfig` for more info.
+        scaling_config: The configuration for how to scale data parallel training.
+            ``num_workers`` determines how many Python processes are used for training,
+            and ``use_gpu`` determines whether or not each process should use GPUs.
+            See :class:`~ray.train.ScalingConfig` for more info.
+        run_config: The configuration for the execution of the training run.
+            See :class:`~ray.train.RunConfig` for more info.
+        datasets: The Ray Datasets to use for training and validation.
+        dataset_config: The configuration for ingesting the input ``datasets``.
+            By default, all the Ray Datasets are split equally across workers.
+            See :class:`~ray.train.DataConfig` for more details.
+        resume_from_checkpoint: A checkpoint to resume training from.
+            This checkpoint can be accessed from within ``train_loop_per_worker``
+            by calling ``ray.train.get_checkpoint()``.
+        metadata: Dict that should be made available via
+            `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
+            for checkpoints saved from this Trainer. Must be JSON-serializable.
+        label_column: [Deprecated] Name of the label column. A column with this name
+            must be present in the training dataset.
+        params: [Deprecated] LightGBM training parameters.
+            Refer to `LightGBM documentation <https://lightgbm.readthedocs.io/>`_
+            for a list of possible parameters.
+        num_boost_round: [Deprecated] Target number of boosting iterations (trees in the model).
+            Note that unlike in ``lightgbm.train``, this is the target number
+            of trees, meaning that if you set ``num_boost_round=10`` and pass a model
+            that has already been trained for 5 iterations, it will be trained for 5
+            iterations more, instead of 10 more.
+        **train_kwargs: [Deprecated] Additional kwargs passed to ``lightgbm.train()`` function.
+    """
+
+    _handles_checkpoint_freq = True
+    _handles_checkpoint_at_end = True
+
+    def __init__(
+        self,
+        train_loop_per_worker: Optional[
+            Union[Callable[[], None], Callable[[Dict], None]]
+        ] = None,
+        *,
+        train_loop_config: Optional[Dict] = None,
+        lightgbm_config: Optional[LightGBMConfig] = None,
+        scaling_config: Optional[ray.train.ScalingConfig] = None,
+        run_config: Optional[ray.train.RunConfig] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        dataset_config: Optional[ray.train.DataConfig] = None,
+        resume_from_checkpoint: Optional[Checkpoint] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        # TODO: [Deprecated] Legacy LightGBMTrainer API
+        label_column: Optional[str] = None,
+        params: Optional[Dict[str, Any]] = None,
+        num_boost_round: Optional[int] = None,
+        **train_kwargs,
+    ):
+        # TODO: [Deprecated] Legacy LightGBMTrainer API
+        legacy_api = train_loop_per_worker is None
+        if legacy_api:
+            train_loop_per_worker = self._get_legacy_train_fn_per_worker(
+                lightgbm_train_kwargs=train_kwargs,
+                run_config=run_config,
+                label_column=label_column,
+                num_boost_round=num_boost_round,
+                datasets=datasets,
+            )
+            train_loop_config = params or {}
+        elif train_kwargs:
+            _log_deprecation_warning(
+                "Passing `lightgbm.train` kwargs to `LightGBMTrainer` is deprecated. "
+                f"Got kwargs: {train_kwargs.keys()}\n"
+                "In your training function, you can call `lightgbm.train(**kwargs)` "
+                "with arbitrary arguments. "
+                f"{LEGACY_LIGHTGBM_TRAINER_DEPRECATION_MESSAGE}"
+            )
+
+        super(LightGBMTrainer, self).__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            train_loop_config=train_loop_config,
+            lightgbm_config=lightgbm_config,
+            scaling_config=scaling_config,
+            run_config=run_config,
+            datasets=datasets,
+            dataset_config=dataset_config,
+            resume_from_checkpoint=resume_from_checkpoint,
+            metadata=metadata,
+        )
+
+    def _get_legacy_train_fn_per_worker(
+        self,
+        lightgbm_train_kwargs: Dict,
+        run_config: Optional[ray.train.RunConfig],
+        datasets: Optional[Dict[str, GenDataset]],
+        label_column: Optional[str],
+        num_boost_round: Optional[int],
+    ) -> Callable[[Dict], None]:
+        """Get the training function for the legacy LightGBMTrainer API."""
+
+        datasets = datasets or {}
+        if not datasets.get(TRAIN_DATASET_KEY):
+            raise ValueError(
+                "`datasets` must be provided for the LightGBMTrainer API "
+                "if `train_loop_per_worker` is not provided. "
+                "This dict must contain the training dataset under the "
+                f"key: '{TRAIN_DATASET_KEY}'. "
+                f"Got keys: {list(datasets.keys())}"
+            )
+        if not label_column:
+            raise ValueError(
+                "`label_column` must be provided for the LightGBMTrainer API "
+                "if `train_loop_per_worker` is not provided. "
+                "This is the column name of the label in the dataset."
+            )
+
+        num_boost_round = num_boost_round or 10
+
+        _log_deprecation_warning(LEGACY_LIGHTGBM_TRAINER_DEPRECATION_MESSAGE)
+
+        # Initialize a default Ray Train metrics/checkpoint reporting callback if needed
+        callbacks = lightgbm_train_kwargs.get("callbacks", [])
+        user_supplied_callback = any(
+            isinstance(callback, RayTrainReportCallback) for callback in callbacks
+        )
+        callback_kwargs = {}
+        if run_config:
+            checkpoint_frequency = run_config.checkpoint_config.checkpoint_frequency
+            checkpoint_at_end = run_config.checkpoint_config.checkpoint_at_end
+
+            callback_kwargs["frequency"] = checkpoint_frequency
+            # Default `checkpoint_at_end=True` unless the user explicitly sets it.
+            callback_kwargs["checkpoint_at_end"] = (
+                checkpoint_at_end if checkpoint_at_end is not None else True
+            )
+
+        if not user_supplied_callback:
+            callbacks.append(RayTrainReportCallback(**callback_kwargs))
+        lightgbm_train_kwargs["callbacks"] = callbacks
+
+        train_fn_per_worker = partial(
+            _lightgbm_train_fn_per_worker,
+            label_column=label_column,
+            num_boost_round=num_boost_round,
+            dataset_keys=set(datasets),
+            lightgbm_train_kwargs=lightgbm_train_kwargs,
+        )
+        return train_fn_per_worker
+
+    @classmethod
+    def get_model(
+        cls,
+        checkpoint: Checkpoint,
+    ) -> lightgbm.Booster:
+        """Retrieve the LightGBM model stored in this checkpoint."""
+        return RayTrainReportCallback.get_model(checkpoint)

python/ray/train/lightgbm/v2.py

@@ -0,0 +1,131 @@
+import logging
+from typing import Any, Callable, Dict, Optional, Union
+
+import ray.train
+from ray.train import Checkpoint
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.lightgbm.config import LightGBMConfig, get_network_params  # noqa: F401
+from ray.train.trainer import GenDataset
+
+logger = logging.getLogger(__name__)
+
+
+class LightGBMTrainer(DataParallelTrainer):
+    """A Trainer for distributed data-parallel LightGBM training.
+
+    Example:
+
+        .. testcode::
+            :skipif: True
+
+            import lightgbm as lgb
+
+            import ray.data
+            import ray.train
+            from ray.train.lightgbm import (
+                LightGBMTrainer,
+                RayTrainReportCallback,
+                normalize_pandas_for_lightgbm,
+            )
+
+
+            def train_fn_per_worker(config: dict):
+                # (Optional) Add logic to resume training state from a checkpoint.
+                # ray.train.get_checkpoint()
+
+                # 1. Get the dataset shard for the worker and convert to a `lgb.Dataset`
+                train_ds_iter, eval_ds_iter = (
+                    ray.train.get_dataset_shard("train"),
+                    ray.train.get_dataset_shard("validation"),
+                )
+                train_ds, eval_ds = train_ds_iter.materialize(), eval_ds_iter.materialize()
+                train_df = normalize_pandas_for_lightgbm(train_ds.to_pandas())
+                eval_df = normalize_pandas_for_lightgbm(eval_ds.to_pandas())
+                train_X, train_y = train_df.drop("y", axis=1), train_df["y"]
+                eval_X, eval_y = eval_df.drop("y", axis=1), eval_df["y"]
+
+                train_set = lgb.Dataset(train_X, label=train_y)
+                eval_set = lgb.Dataset(eval_X, label=eval_y)
+
+                # 2. Run distributed data-parallel training.
+                # `get_network_params` sets up the necessary configurations for LightGBM
+                # to set up the data parallel training worker group on your Ray cluster.
+                params = {
+                    "objective": "regression",
+                    # Adding the line below is the only change needed
+                    # for your `lgb.train` call!
+                    **ray.train.lightgbm.get_network_params(),
+                }
+                lgb.train(
+                    params,
+                    train_set,
+                    valid_sets=[eval_set],
+                    valid_names=["eval"],
+                    callbacks=[RayTrainReportCallback()],
+                )
+
+            train_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
+            eval_ds = ray.data.from_items(
+                [{"x": x, "y": x + 1} for x in range(32, 32 + 16)]
+            )
+            trainer = LightGBMTrainer(
+                train_fn_per_worker,
+                datasets={"train": train_ds, "validation": eval_ds},
+                scaling_config=ray.train.ScalingConfig(num_workers=4),
+            )
+            result = trainer.fit()
+            booster = RayTrainReportCallback.get_model(result.checkpoint)
+
+    Args:
+        train_loop_per_worker: The training function to execute on each worker.
+            This function can either take in zero arguments or a single ``Dict``
+            argument which is set by defining ``train_loop_config``.
+            Within this function you can use any of the
+            :ref:`Ray Train Loop utilities <train-loop-api>`.
+        train_loop_config: A configuration ``Dict`` to pass in as an argument to
+            ``train_loop_per_worker``.
+            This is typically used for specifying hyperparameters.
+        lightgbm_config: The configuration for setting up the distributed lightgbm
+            backend. See :class:`~ray.train.lightgbm.LightGBMConfig` for more info.
+        scaling_config: The configuration for how to scale data parallel training.
+            ``num_workers`` determines how many Python processes are used for training,
+            and ``use_gpu`` determines whether or not each process should use GPUs.
+            See :class:`~ray.train.ScalingConfig` for more info.
+        run_config: The configuration for the execution of the training run.
+            See :class:`~ray.train.RunConfig` for more info.
+        datasets: The Ray Datasets to use for training and validation.
+        dataset_config: The configuration for ingesting the input ``datasets``.
+            By default, all the Ray Dataset are split equally across workers.
+            See :class:`~ray.train.DataConfig` for more details.
+        metadata: Dict that should be made available via
+            `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
+            for checkpoints saved from this Trainer. Must be JSON-serializable.
+        resume_from_checkpoint: A checkpoint to resume training from.
+            This checkpoint can be accessed from within ``train_loop_per_worker``
+            by calling ``ray.train.get_checkpoint()``.
+    """
+
+    def __init__(
+        self,
+        train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]],
+        *,
+        train_loop_config: Optional[Dict] = None,
+        lightgbm_config: Optional[LightGBMConfig] = None,
+        scaling_config: Optional[ray.train.ScalingConfig] = None,
+        run_config: Optional[ray.train.RunConfig] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        dataset_config: Optional[ray.train.DataConfig] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        resume_from_checkpoint: Optional[Checkpoint] = None,
+    ):
+        super(LightGBMTrainer, self).__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            train_loop_config=train_loop_config,
+            backend_config=lightgbm_config or LightGBMConfig(),
+            scaling_config=scaling_config,
+            dataset_config=dataset_config,
+            run_config=run_config,
+            datasets=datasets,
+            resume_from_checkpoint=resume_from_checkpoint,
+            metadata=metadata,
+        )

python/ray/train/lightning/__init__.py

@@ -0,0 +1,33 @@
+# isort: off
+try:
+    import lightning  # noqa: F401
+except ModuleNotFoundError:
+    try:
+        import pytorch_lightning  # noqa: F401
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "PyTorch Lightning isn't installed. To install PyTorch Lightning, "
+            "please run 'pip install lightning'"
+        )
+# isort: on
+
+from ray.train.lightning._lightning_utils import (
+    RayDDPStrategy,
+    RayDeepSpeedStrategy,
+    RayFSDPStrategy,
+    RayLightningEnvironment,
+    RayTrainReportCallback,
+    prepare_trainer,
+)
+
+__all__ = [
+    "prepare_trainer",
+    "RayDDPStrategy",
+    "RayFSDPStrategy",
+    "RayDeepSpeedStrategy",
+    "RayLightningEnvironment",
+    "RayTrainReportCallback",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/lightning/_lightning_utils.py

@@ -0,0 +1,353 @@
+import logging
+import os
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Dict
+
+import torch
+from packaging.version import Version
+
+import ray
+import ray.train
+from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
+from ray.train import Checkpoint
+from ray.train.v2._internal.constants import is_v2_enabled
+from ray.util import PublicAPI
+
+
+def import_lightning():  # noqa: F402
+    try:
+        import lightning.pytorch as pl
+    except ModuleNotFoundError:
+        import pytorch_lightning as pl
+    return pl
+
+
+pl = import_lightning()
+
+if is_v2_enabled():
+    from ray.train.v2.api.report_config import CheckpointUploadMode
+
+_LIGHTNING_GREATER_EQUAL_2_0 = Version(pl.__version__) >= Version("2.0.0")
+_LIGHTNING_LESS_THAN_2_1 = Version(pl.__version__) < Version("2.1.0")
+_TORCH_GREATER_EQUAL_1_12 = Version(torch.__version__) >= Version("1.12.0")
+_TORCH_FSDP_AVAILABLE = _TORCH_GREATER_EQUAL_1_12 and torch.distributed.is_available()
+
+try:
+    from lightning.pytorch.plugins.environments import LightningEnvironment
+except ModuleNotFoundError:
+    from pytorch_lightning.plugins.environments import LightningEnvironment
+
+if _LIGHTNING_GREATER_EQUAL_2_0:
+    FSDPStrategy = pl.strategies.FSDPStrategy
+else:
+    FSDPStrategy = pl.strategies.DDPFullyShardedStrategy
+
+if _TORCH_FSDP_AVAILABLE:
+    from torch.distributed.fsdp import (
+        FullStateDictConfig,
+        FullyShardedDataParallel,
+        StateDictType,
+    )
+
+
+logger = logging.getLogger(__name__)
+
+LIGHTNING_REPORT_STAGE_KEY = "_report_on"
+
+
+@PublicAPI(stability="beta")
+class RayDDPStrategy(pl.strategies.DDPStrategy):
+    """Subclass of DDPStrategy to ensure compatibility with Ray orchestration.
+
+    For a full list of initialization arguments, please refer to:
+    https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.DDPStrategy.html
+
+    Note that `process_group_backend`, `timeout`, and `start_method` are disabled here,
+    please specify these arguments in :class:`~ray.train.torch.TorchConfig` instead.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYDDPSTRATEGY, "1")
+
+    @property
+    def root_device(self) -> torch.device:
+        return ray.train.torch.get_device()
+
+    @property
+    def distributed_sampler_kwargs(self) -> Dict[str, Any]:
+        return dict(
+            num_replicas=self.world_size,
+            rank=self.global_rank,
+        )
+
+
+@PublicAPI(stability="beta")
+class RayFSDPStrategy(FSDPStrategy):  # noqa: F821
+    """Subclass of FSDPStrategy to ensure compatibility with Ray orchestration.
+
+    For a full list of initialization arguments, please refer to:
+    https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.FSDPStrategy.html
+
+    .. note::
+        It is recommended to upgrade `lightning>=2.1` or above when using FSDP
+        with Lightning, since Lightning starts to natively support `state_dict_type`,
+        `sharding_strategy`, `auto_wrap_policy` and other FSDP configurations from 2.1.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYFSDPSTRATEGY, "1")
+
+    @property
+    def root_device(self) -> torch.device:
+        return ray.train.torch.get_device()
+
+    @property
+    def distributed_sampler_kwargs(self) -> Dict[str, Any]:
+        return dict(
+            num_replicas=self.world_size,
+            rank=self.global_rank,
+        )
+
+    def lightning_module_state_dict(self) -> Dict[str, Any]:
+        """Gathers the full state dict to rank 0 on CPU.
+
+        FSDP checkpointing is broken in Lightning 2.0.x. This subclass patches the
+        behavior to perform a full state dict checkpointing, gathering the checkpoint
+        shards on rank 0 CPU. Upgrade to `lightning>=2.1` to do sharded state dict
+        checkpointing.
+
+        See the note in the class docstring for more details.
+        """
+
+        assert self.model is not None, "Failed to get the state dict for a None model!"
+
+        if (
+            _TORCH_FSDP_AVAILABLE
+            and _LIGHTNING_GREATER_EQUAL_2_0
+            and _LIGHTNING_LESS_THAN_2_1
+        ):
+            with FullyShardedDataParallel.state_dict_type(
+                module=self.model,
+                state_dict_type=StateDictType.FULL_STATE_DICT,
+                state_dict_config=FullStateDictConfig(
+                    offload_to_cpu=True, rank0_only=True
+                ),
+            ):
+                state_dict = self.model.state_dict()
+
+                ckpt_state_dict = {}
+                prefix_len = len("_forward_module.")
+                for k, v in state_dict.items():
+                    if k.startswith("_forward_module."):
+                        non_prefixed_key = k[prefix_len:]
+                        ckpt_state_dict[non_prefixed_key] = v
+                    else:
+                        ckpt_state_dict[k] = v
+                return ckpt_state_dict
+        else:
+            # Otherwise Lightning uses Fairscale FSDP, no need to unshard by ourself.
+            return super().lightning_module_state_dict()
+
+
+@PublicAPI(stability="beta")
+class RayDeepSpeedStrategy(pl.strategies.DeepSpeedStrategy):
+    """Subclass of DeepSpeedStrategy to ensure compatibility with Ray orchestration.
+
+    For a full list of initialization arguments, please refer to:
+    https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.DeepSpeedStrategy.html
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYDEEPSPEEDSTRATEGY, "1")
+
+    @property
+    def root_device(self) -> torch.device:
+        return ray.train.torch.get_device()
+
+    @property
+    def distributed_sampler_kwargs(self) -> Dict[str, Any]:
+        return dict(
+            num_replicas=self.world_size,
+            rank=self.global_rank,
+        )
+
+
+@PublicAPI(stability="beta")
+class RayLightningEnvironment(LightningEnvironment):  # noqa: F821
+    """Setup Lightning DDP training environment for Ray cluster."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYLIGHTNINGENVIRONMENT, "1")
+
+    def world_size(self) -> int:
+        return ray.train.get_context().get_world_size()
+
+    def global_rank(self) -> int:
+        return ray.train.get_context().get_world_rank()
+
+    def local_rank(self) -> int:
+        return ray.train.get_context().get_local_rank()
+
+    def node_rank(self) -> int:
+        return ray.train.get_context().get_node_rank()
+
+    def set_world_size(self, size: int) -> None:
+        # Disable it since `world_size()` directly returns data from Train context.
+        pass
+
+    def set_global_rank(self, rank: int) -> None:
+        # Disable it since `global_rank()` directly returns data from Train.
+        pass
+
+    def teardown(self):
+        pass
+
+
+@PublicAPI(stability="beta")
+def prepare_trainer(trainer: pl.Trainer) -> pl.Trainer:
+    """Prepare the PyTorch Lightning Trainer for distributed execution."""
+
+    # Check strategy class
+    valid_strategy_class = [RayDDPStrategy, RayFSDPStrategy, RayDeepSpeedStrategy]
+
+    if not any(isinstance(trainer.strategy, cls) for cls in valid_strategy_class):
+        raise RuntimeError(
+            f"Invalid strategy class: {type(trainer.strategy)}. To use "
+            "PyTorch Lightning with Ray, the strategy object should be one of "
+            f"{[cls.__name__ for cls in valid_strategy_class]} class "
+            "or its subclass."
+        )
+
+    # Check cluster environment
+    cluster_environment = getattr(trainer.strategy, "cluster_environment", None)
+    if cluster_environment and not isinstance(
+        cluster_environment, RayLightningEnvironment
+    ):
+        raise RuntimeError(
+            "Invalid cluster environment plugin. The expected class is"
+            "`ray.train.lightning.RayLightningEnvironment` "
+            f"but got {type(cluster_environment)}!"
+        )
+
+    record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_PREPARE_TRAINER, "1")
+    return trainer
+
+
+@PublicAPI(stability="beta")
+class RayTrainReportCallback(pl.callbacks.Callback):
+    """A simple callback that reports checkpoints to Ray on train epoch end.
+
+    This callback is a subclass of `lightning.pytorch.callbacks.Callback
+    <https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.Callback.html#lightning.pytorch.callbacks.Callback>`_.
+
+    It fetches the latest `trainer.callback_metrics` and reports together with
+    the checkpoint on each training epoch end.
+
+    Checkpoints will be saved in the following structure::
+
+        checkpoint_00000*/      Ray Train Checkpoint
+        └─ checkpoint.ckpt      PyTorch Lightning Checkpoint
+
+    You can also provide the following arguments to the callback:
+
+    - checkpoint_upload_mode: The manner in which to upload the checkpoint.
+        See :ref:`Checkpoint upload modes <train-checkpoint-upload-modes>` for more details.
+    - validation: Whether to asynchronously validate the checkpoint.
+        See :ref:`Validating checkpoints asynchronously <train-validating-checkpoints>` for more details.
+
+    For customized reporting and checkpointing logic, implement your own
+    `lightning.pytorch.callbacks.Callback` following this user
+    guide: :ref:`Saving and Loading Checkpoints <train-dl-saving-checkpoints>`.
+    """
+
+    CHECKPOINT_NAME = "checkpoint.ckpt"
+
+    def __init__(
+        self,
+        # v2-only arguments
+        checkpoint_upload_mode=None,
+        validation=False,
+    ) -> None:
+        super().__init__()
+        self.checkpoint_upload_mode = checkpoint_upload_mode
+        self.validation = validation
+
+        if is_v2_enabled():
+            if checkpoint_upload_mode is None:
+                self.checkpoint_upload_mode = CheckpointUploadMode.SYNC
+        else:
+            if checkpoint_upload_mode is not None:
+                raise ValueError(
+                    "`checkpoint_upload_mode` is only supported in Ray Train v2. "
+                    "To enable it, please set `RAY_TRAIN_V2_ENABLED=1`."
+                )
+            if validation:
+                raise ValueError(
+                    "`validation` is only supported in Ray Train v2. "
+                    "To enable it, please set `RAY_TRAIN_V2_ENABLED=1`."
+                )
+
+        job_id = ray.get_runtime_context().get_job_id()
+        experiment_name = ray.train.get_context().get_experiment_name()
+
+        self.tmpdir_prefix = Path(
+            tempfile.gettempdir(),
+            f"lightning_checkpoints-job_id={job_id}-name={experiment_name}-world_rank={ray.train.get_context().get_world_rank()}",
+        ).as_posix()
+        if os.path.isdir(self.tmpdir_prefix):
+            shutil.rmtree(self.tmpdir_prefix)
+
+        record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYTRAINREPORTCALLBACK, "1")
+
+    def on_train_epoch_end(self, trainer, pl_module) -> None:
+        # Creates a checkpoint dir with fixed name
+        tmpdir = Path(self.tmpdir_prefix, str(trainer.current_epoch)).as_posix()
+        os.makedirs(tmpdir, exist_ok=True)
+
+        # Fetch metrics
+        metrics = trainer.callback_metrics
+        metrics = {k: v.item() for k, v in metrics.items()}
+
+        # (Optional) Add customized metrics
+        metrics["epoch"] = trainer.current_epoch
+        metrics["step"] = trainer.global_step
+
+        # Save checkpoint to local
+        ckpt_path = Path(tmpdir, self.CHECKPOINT_NAME).as_posix()
+        # TODO: with CheckpointUploadMode.ASYNC, this does cpu -> disk synchronously
+        # and disk -> remote asynchronously. We can add a new CheckpointIO class to do
+        # cpu -> remote asynchronously and a checkpoint_upload_fn to wait for it.
+        trainer.save_checkpoint(ckpt_path, weights_only=False)
+
+        # Report to train session
+        checkpoint = Checkpoint.from_directory(tmpdir)
+        if is_v2_enabled():
+            ray.train.report(
+                metrics=metrics,
+                checkpoint=checkpoint,
+                checkpoint_upload_mode=self.checkpoint_upload_mode,
+                validation=self.validation,
+            )
+        else:
+            ray.train.report(metrics=metrics, checkpoint=checkpoint)
+
+        # Add a barrier to ensure all workers finished reporting here
+        trainer.strategy.barrier()
+
+        # With CheckpointUploadMode.ASYNC, the upload may still be in progress
+        # after report() returns. Let ray.train.report delete_local_checkpoint_after_upload
+        # handle cleanup instead.
+        if is_v2_enabled() and self.checkpoint_upload_mode is not None:
+            # Check here because CheckpointUploadMode is only imported when is_v2_enabled() is True
+            if (
+                self.checkpoint_upload_mode.default_delete_local_checkpoint_after_upload()
+            ):
+                return
+
+        shutil.rmtree(tmpdir)

python/ray/train/lint/README.md

@@ -0,0 +1,29 @@
+## Ray Train Circular Import Linter
+
+Ray Train functionality is overrided or "patched" by functionality from other directories. For instance, Ray Train is patched by functionality from Ray Train v2 when `RAY_TRAIN_V2_ENABLED=1`, making Ray Train dependent on Ray Train v2. In turn, the patching directory often imports functionality from the "base" Ray Train directory (`ray/python/ray/train`), resulting in a circular dependency. The Ray Train Circular Import Linter takes a patching directory, `patch_dir`, and detects circular imports between it and the base Ray Train directory- displaying violations to users via pre-commit.
+
+### The Problem:
+
+When there is a circular dependency present in the codebase, import errors are triggered by importing directly from the patching directory. For example, directly importing the v2 TensorflowTrainer with `from ray.train.v2.tensorflow.tensorflow_trainer import TensorflowTrainer` rather than relying on the conventional routing logic via `from ray.train.tensorflow import TensorflowTrainer` results in a `ImportError: cannot import name TensorflowTrainer from partially initialized ray.train.v2.tensorflow.tensorflow_trainer`. This happens in the case of user API misuse or during the deserialization of the train function on train worker setup. The following image depicts the error path of such erroneous imports.
+
+![ErrorPath](./images/ErrorPath.png)
+
+### The Fix:
+
+To make Ray Train resilient to such erroneous imports, this linter proactively detects circular imports and specifies how to fix it. The fix perscribed by the linter prevents import errors by importing the base Train packages early within in the patching directory. In the below example, the previously depicted circular import is resolved by the linter's suggested fix to import `ray.train.foo` early in `ray.train.v2.foo`.
+
+![SuccessPath](./images/SuccessPath.png)
+
+The key observation is that the fix redirects the import path of `from ray.train.v2.foo import foo_v2` so that the base Train init file (e.g. `ray.train.foo.__init__.py`) runs before the patching file (e.g. `ray.train.v2.foo.py`), avoiding the error in the previous example.
+
+### Linter Specification
+
+The linter implements an `ast.NodeVisitor` to parse imports within the base Train directory and the patching directory to detect circular imports. The below example depicts two circular imports that would be detected by the linter originating from a `ray.train.common.__init__.py` file.
+
+![Linter](./images/Linter.png)
+
+The linter parses all `__init__.py` files in the base Train directory and collects their imports. For each patching import (e.g. `from ray.train.v2.foo import foo_v2`), the linter will also collect the imports in the patching file (e.g. `ray.train.v2.foo.py`) and if any of these imports point back to the same base Train file (e.g. `ray.train.common.__init__.py`), a violation is detected.
+
+However, notice from the diagram that the linter also detects violations in the case of reexports. If the base Train file points to a patching package file (e.g. `ray.train.v2.bar.__init__.py`), the linter will also collect the imports of the referenced implementation file (e.g. `ray.train.v2.bar.bar_impl.py`) to search for a violation.
+
+That said, in both cases, if the linter finds that the base Train file is imported early in the patching package file (e.g. `ray.train.common` is imported in `ray.train.v2.foo.__init__.py`/`ray.train.v2.bar.__init__.py`), then the violation will be suppressed. Otherwise, this is the fix that will be reccommended by the linter.

python/ray/train/lint/check_circular_imports.py

@@ -0,0 +1,393 @@
+import argparse
+import ast
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Set, Tuple
+
+TRAIN_PACKAGES = set()
+
+
+def find_train_packages(base_train_dir: Path, patch_train_dir: Path) -> None:
+    """
+    Find and initialize the global TRAIN_PACKAGES with all train package names from base_train_dir
+    and patch_train_dir combined.
+    """
+    global TRAIN_PACKAGES
+    TRAIN_PACKAGES = set()
+
+    # Collect all packages under both base and patch train dirs
+    package_files = list(base_train_dir.rglob("__init__.py")) + list(
+        patch_train_dir.rglob("__init__.py")
+    )
+    base_dir = get_base_dir()
+    for init_file in package_files:
+        relative_path = init_file.relative_to(base_dir)
+        dotted_module = str(relative_path.parent).replace("/", ".")
+        TRAIN_PACKAGES.add(dotted_module)
+
+
+def is_train_package(module_str: str) -> bool:
+    return module_str in TRAIN_PACKAGES
+
+
+def get_base_dir() -> Path:
+    """Return the filesystem path to the ray python directory."""
+    current_file_path = Path(__file__).resolve()
+    package_dir = current_file_path.parents[3]
+    return package_dir
+
+
+def get_base_train_dir() -> Path:
+    """Return the filesystem path to the ray train directory."""
+    return get_base_dir() / "ray/train"
+
+
+def does_overlap(main_module: str, module: str) -> bool:
+    """Checks if the init file of module is on the import path of main_module"""
+    return main_module.startswith(f"{module}.") or main_module == module
+
+
+class Import:
+    """
+    Represents an import statement.
+    For example, 'from X import A, B' has module 'X' and names ['A', 'B'].
+    Also supports 'import X'.
+    """
+
+    def __init__(
+        self, module: str, names: List[str] = None, is_package: bool = False
+    ) -> None:
+        self.is_package = is_package
+        self.module = module
+        self.names = names if names else []
+
+
+class ImportCollector(ast.NodeVisitor):
+    """
+    An AST node visitor that collects all module-level imports from a Python source file.
+    It traverses the AST and records module-level import statements (`import ...` and `from ... import ...`) that are not
+    inside function or class definitions, and that are not guarded by `if TYPE_CHECKING` or `if typing.TYPE_CHECKING`
+    blocks.
+    """
+
+    def __init__(self, module_name: str, is_package: bool) -> None:
+        self._module_name = module_name
+        self._is_package = is_package
+        self.imports: Set[Import] = set()
+        self.type_checking_imported = False
+
+    # --- private helpers ---
+
+    def _is_type_checking_test(self, expr: ast.AST) -> bool:
+        """Return True for `if TYPE_CHECKING` or `if typing.TYPE_CHECKING`."""
+
+        if (
+            self.type_checking_imported
+            and isinstance(expr, ast.Name)
+            and expr.id == "TYPE_CHECKING"
+        ):
+            return True
+        elif (
+            isinstance(expr, ast.Attribute)
+            and isinstance(expr.value, ast.Name)
+            and expr.value.id == "typing"
+            and expr.attr == "TYPE_CHECKING"
+        ):
+            return True
+
+        return False
+
+    def _get_package_parts(self) -> List[str]:
+        parts = self._module_name.split(".")
+        return parts if self._is_package else parts[:-1]
+
+    def _to_absolute_module(
+        self, level: int, module_str: Optional[str]
+    ) -> Optional[str]:
+        """Construct the absolute module string from a relative import."""
+        # Absolute import
+        if level == 0:
+            return module_str
+
+        package_parts = self._get_package_parts()
+
+        # If the relative import is out of bounds
+        if level - 1 > len(package_parts):
+            return None
+
+        # Base parts based on the level
+        base_module_parts = (
+            package_parts if level == 1 else package_parts[: -(level - 1)]
+        )
+
+        # Construct absolute module string
+        abs_module_parts = (
+            base_module_parts + module_str.split(".")
+            if module_str
+            else base_module_parts
+        )
+        return ".".join(abs_module_parts)
+
+    # --- parsing functions ---
+
+    def visit_If(self, node: ast.If) -> None:
+        # If the test is not TYPE_CHECKING, visit statement body
+        if not self._is_type_checking_test(node.test):
+            for stmt in node.body:
+                self.visit(stmt)
+
+        # Also visit conditional branches
+        for stmt in node.orelse:
+            self.visit(stmt)
+
+    def visit_Import(self, node: ast.Import) -> None:
+        for alias in node.names:
+            if alias.name:
+                self.imports.add(
+                    Import(module=alias.name, is_package=is_train_package(alias.name))
+                )
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        import_str = self._to_absolute_module(node.level or 0, node.module)
+        if not import_str:
+            return
+
+        names = [alias.name for alias in node.names]
+        self.imports.add(
+            Import(
+                module=import_str, is_package=is_train_package(import_str), names=names
+            )
+        )
+        if "TYPE_CHECKING" in names and import_str == "typing":
+            self.type_checking_imported = True
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        # Skip function contents
+        return
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
+        # Skip function contents
+        return
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        # Skip class contents
+        return
+
+
+def collect_imports(
+    module_name: str, is_package: bool, source_text: str
+) -> Set[Import]:
+    try:
+        tree = ast.parse(source_text)
+    except SyntaxError:
+        print(f"Warning: Failed to parse {module_name} for circular imports")
+        return set()
+    collector = ImportCollector(module_name, is_package)
+    collector.visit(tree)
+    return collector.imports
+
+
+def to_module_name_and_is_package(py_file: Path) -> Tuple[str, bool]:
+    """
+    Convert a Python file path to its corresponding module name and determine if it is a package.
+
+    Args:
+        py_file: The path to the Python file.
+
+    Returns:
+        Tuple[str, bool]: A tuple containing the module name as a string and a boolean indicating
+                          whether the module is a package (True if it is an __init__.py file).
+    """
+    file_path = py_file.relative_to(get_base_dir())
+    module_path = file_path.with_suffix("")
+    module_parts = module_path.parts
+    is_package = module_parts[-1] == "__init__"
+    if is_package:
+        module_parts = module_parts[:-1]
+    module_str = ".".join(module_parts)
+    return module_str, is_package
+
+
+def get_file_module_imports(
+    files: List[Path], module_match_string: Optional[str] = None
+) -> Dict[str, List[Import]]:
+    """
+    Collect and return the module-level imports for a list of Python files.
+
+    Args:
+        files: A list of Path objects representing Python files to analyze.
+        module_match_string: An optional string to filter imports. Only imports
+                             containing this string will be included in the result.
+
+    Returns:
+        A dictionary mapping module names to a list of their import statements.
+        The module names are derived from the file paths, and the import statements
+        are filtered based on the optional module_match_string.
+    """
+    module_imports: Dict[str, List[Import]] = {}
+
+    # Collect the imports for each python file
+    for py_file in files:
+        try:
+            module_name, is_package = to_module_name_and_is_package(py_file)
+            src = py_file.read_text(encoding="utf-8", errors="ignore")
+            imports = collect_imports(module_name, is_package, src)
+            module_imports[module_name] = [
+                stmt
+                for stmt in imports
+                if module_match_string is None or module_match_string in stmt.module
+            ]
+        except Exception:
+            continue
+    return module_imports
+
+
+def convert_to_file_paths(imports: List[Import]) -> List[Path]:
+    """
+    Convert a list of import strings to a list of file paths.
+
+    Args:
+        imports: A list of Import objects
+
+    Returns:
+        A list of file paths.
+    """
+    base_dir = get_base_dir()
+    file_paths = []
+    for imp in imports:
+        if imp.is_package:
+            relative_path = imp.module.replace(".", "/") + "/__init__.py"
+        else:
+            relative_path = imp.module.replace(".", "/") + ".py"
+        file_paths.append(base_dir / relative_path)
+    return file_paths
+
+
+def expand_to_include_reexports(import_map: Dict[str, List[Import]]) -> None:
+    """
+    Expands the set of imports for a given import map to include the modules resulting from reexports.
+    So if in the base train module, there is "from x import a, b" and x is a package, then this function
+    will explore the __init__.py of x and include the modules a and b were reexported from in the import map.
+    """
+    for module, base_imports in import_map.items():
+        # Get only the package imports
+        packages = [imp for imp in base_imports if imp.is_package]
+        package_files = convert_to_file_paths(packages)
+        reexports = get_file_module_imports(package_files)
+
+        agg_reexports = []
+        # Filter patch init file imports to those that only contain the right names
+        for base_import in base_imports:
+            if base_import.module in reexports:
+                import_list = reexports[base_import.module]
+                target_reexports = [
+                    imp
+                    for imp in import_list
+                    if set(imp.names) & set(base_import.names)
+                ]
+                agg_reexports.extend(target_reexports)
+
+        # Expand modules to include reexported modules
+        import_map[module].extend(agg_reexports)
+
+
+def check_violations(
+    base_train_patching_imports: Dict[str, List[Import]], patch_dir: Path
+) -> List[str]:
+    """
+    Check for circular import violations between base and patch train modules.
+
+    Args:
+        base_train_patching_imports: A dictionary mapping base train module names to their imports.
+        patch_dir: The directory path containing patch train modules.
+
+    Returns:
+        A list of strings describing any circular import violations found.
+    """
+    violations: List[str] = []
+
+    # Get the imports from the patch train init files
+    patch_train_init_files = list(patch_dir.rglob("__init__.py"))
+    patch_train_init_imports = get_file_module_imports(
+        patch_train_init_files, module_match_string="ray.train"
+    )
+
+    # Expand the imports to include reexports
+    expand_to_include_reexports(base_train_patching_imports)
+
+    # Process each patch train init module for violations
+    for base_train_init_module, imports in base_train_patching_imports.items():
+
+        # Get the imports from the patch train files
+        patch_train_files = convert_to_file_paths(imports)
+        patch_train_file_imports = get_file_module_imports(
+            patch_train_files, module_match_string="ray.train"
+        )
+
+        for patch_module, imports in patch_train_file_imports.items():
+            # Skip if the base train init module is in the import path of the patch module
+            if does_overlap(patch_module, base_train_init_module):
+                continue
+
+            # Skip if the patch train module init file imports the base train init module
+            patch_init_module = (
+                ".".join(patch_module.split(".")[:-1])
+                if not is_train_package(patch_module)
+                else patch_module
+            )
+            patch_init_imports = patch_train_init_imports.get(patch_init_module, [])
+            if any(
+                does_overlap(imp.module, base_train_init_module)
+                for imp in patch_init_imports
+            ):
+                continue
+
+            for patch_import in imports:
+                # If any of those v1 imports go through the init file, then it is a violation
+                if does_overlap(patch_import.module, base_train_init_module):
+                    violations.append(
+                        f"circular-import-train: Circular import between {base_train_init_module} (importing {patch_module}) and {patch_module} (importing {patch_import.module}). Resolve by importing {base_train_init_module} in the __init__.py of {patch_init_module}."
+                    )
+
+    return violations
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--patch_dir",
+        default="ray/train/v2",
+        help="Path to the directory containing patching contents",
+    )
+    args = parser.parse_args()
+
+    # Get train directory paths
+    base_dir = get_base_dir()
+    base_train_dir = get_base_train_dir()
+    patch_train_dir = base_dir / Path(args.patch_dir)
+
+    # Find and save all train packages in global TRAIN_PACKAGES for reference
+    find_train_packages(base_train_dir, patch_train_dir)
+
+    # Collect all base train init files
+    base_train_init_files = [
+        f
+        for f in base_train_dir.rglob("__init__.py")
+        if not f.is_relative_to(patch_train_dir)
+    ]
+
+    # Get the patching imports in the base train init files
+    dotted_module_prefix = str(patch_train_dir.relative_to(base_dir)).replace("/", ".")
+    patching_imports = get_file_module_imports(
+        base_train_init_files, module_match_string=dotted_module_prefix
+    )
+
+    # Collect all violations based off the patching imports
+    violations = check_violations(patching_imports, patch_train_dir)
+    if violations:
+        print("\n".join(violations))
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

python/ray/train/tensorflow/__init__.py

@@ -0,0 +1,30 @@
+# isort: off
+try:
+    import tensorflow as tf  # noqa: F401
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "TensorFlow isn't installed. To install TensorFlow, run 'pip install "
+        "tensorflow'."
+    )
+# isort: on
+
+from ray.train.tensorflow.config import TensorflowConfig
+from ray.train.tensorflow.tensorflow_checkpoint import TensorflowCheckpoint
+from ray.train.tensorflow.tensorflow_trainer import TensorflowTrainer
+from ray.train.tensorflow.train_loop_utils import prepare_dataset_shard
+from ray.train.v2._internal.constants import is_v2_enabled
+
+if is_v2_enabled():
+    from ray.train.v2.tensorflow.tensorflow_trainer import (  # noqa: F811
+        TensorflowTrainer,
+    )
+
+__all__ = [
+    "TensorflowCheckpoint",
+    "TensorflowConfig",
+    "prepare_dataset_shard",
+    "TensorflowTrainer",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/tensorflow/config.py

@@ -0,0 +1,70 @@
+import json
+import logging
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+import ray
+from ray._common.network_utils import build_address
+from ray.train._internal.base_worker_group import BaseWorkerGroup
+from ray.train._internal.utils import get_address_and_port
+from ray.train.backend import Backend, BackendConfig
+from ray.train.v2._internal.util import TrainingFramework
+from ray.util import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="beta")
+@dataclass
+class TensorflowConfig(BackendConfig):
+    @property
+    def backend_cls(self):
+        return _TensorflowBackend
+
+    @property
+    def framework(self):
+        return TrainingFramework.TENSORFLOW
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {}
+
+
+def _setup_tensorflow_environment(worker_addresses: List[str], index: int):
+    """Set up distributed Tensorflow training information.
+
+    This function should be called on each worker.
+
+    Args:
+        worker_addresses: Addresses of all the workers.
+        index: Index (i.e. world rank) of the current worker.
+    """
+    tf_config = {
+        "cluster": {"worker": worker_addresses},
+        "task": {"type": "worker", "index": index},
+    }
+    os.environ["TF_CONFIG"] = json.dumps(tf_config)
+    os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+
+class _TensorflowBackend(Backend):
+    def on_start(self, worker_group: BaseWorkerGroup, backend_config: TensorflowConfig):
+        # Compute URL for initializing distributed setup.
+        def get_url():
+            address, port = get_address_and_port()
+            return build_address(address, port)
+
+        urls = worker_group.execute(get_url)
+
+        # Get setup tasks in order to throw errors on failure.
+        setup_futures = []
+        for i in range(len(worker_group)):
+            setup_futures.append(
+                worker_group.execute_single_async(
+                    i,
+                    _setup_tensorflow_environment,
+                    worker_addresses=urls,
+                    index=i,
+                )
+            )
+        ray.get(setup_futures)

python/ray/train/tensorflow/keras.py

@@ -0,0 +1,210 @@
+import shutil
+from abc import abstractmethod
+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 Ray Train'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")
+
+
+class RayReportCallback(_Callback):
+    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 _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
+
+    @abstractmethod
+    def _save_and_report_checkpoint(
+        self, metrics: Dict, checkpoint: TensorflowCheckpoint
+    ):
+        """Save checkpoint and report metrics corresonding to this checkpoint."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _report_metrics(self, metrics: Dict):
+        """Report metrics."""
+        raise NotImplementedError
+
+    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)
+            self._save_and_report_checkpoint(metrics, checkpoint)
+            # Clean up temporary checkpoint
+            shutil.rmtree(checkpoint.path, ignore_errors=True)
+        else:
+            self._report_metrics(metrics)
+
+
+@PublicAPI(stability="alpha")
+class ReportCheckpointCallback(RayReportCallback):
+    """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:
+        .. testcode:: 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:
+        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.
+        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".
+        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".
+    """
+
+    def _save_and_report_checkpoint(
+        self, metrics: Dict, checkpoint: TensorflowCheckpoint
+    ):
+        """Save checkpoint and report metrics corresonding to this checkpoint."""
+        ray.train.report(metrics, checkpoint=checkpoint)
+
+    def _report_metrics(self, metrics: Dict):
+        """Report metrics."""
+        ray.train.report(metrics, checkpoint=None)

python/ray/train/tensorflow/tensorflow_checkpoint.py

@@ -0,0 +1,155 @@
+import os
+import shutil
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+import tensorflow as tf
+from tensorflow import keras
+
+from ray.train._internal.framework_checkpoint import FrameworkCheckpoint
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.preprocessor import Preprocessor
+
+
+@PublicAPI(stability="beta")
+class TensorflowCheckpoint(FrameworkCheckpoint):
+    """A :py:class:`~ray.train.Checkpoint` with TensorFlow-specific functionality."""
+
+    MODEL_FILENAME_KEY = "_model_filename"
+
+    @classmethod
+    def from_model(
+        cls,
+        model: keras.Model,
+        *,
+        preprocessor: Optional["Preprocessor"] = None,
+    ) -> "TensorflowCheckpoint":
+        """Create a :py:class:`~ray.train.Checkpoint` that stores a Keras model.
+
+        The checkpoint created with this method needs to be paired with
+        `model` when used.
+
+        Args:
+            model: The Keras model, whose weights are stored in the checkpoint.
+            preprocessor: A fitted preprocessor to be applied before inference.
+
+        Returns:
+            A :py:class:`TensorflowCheckpoint` containing the specified model.
+
+        Examples:
+
+            .. testcode::
+
+                from ray.train.tensorflow import TensorflowCheckpoint
+                import tensorflow as tf
+
+                model = tf.keras.applications.resnet.ResNet101()
+                checkpoint = TensorflowCheckpoint.from_model(model)
+
+            .. testoutput::
+                :options: +MOCK
+                :hide:
+
+                ...  # Model may or may not be downloaded
+
+        """
+        tempdir = tempfile.mkdtemp()
+        filename = "model.keras"
+        model.save(Path(tempdir, filename).as_posix())
+
+        checkpoint = cls.from_directory(tempdir)
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+        checkpoint.update_metadata({cls.MODEL_FILENAME_KEY: filename})
+        return checkpoint
+
+    @classmethod
+    def from_h5(
+        cls, file_path: str, *, preprocessor: Optional["Preprocessor"] = None
+    ) -> "TensorflowCheckpoint":
+        """Create a :py:class:`~ray.train.Checkpoint` that stores a Keras
+        model from H5 format.
+
+        The checkpoint generated by this method contains all the information needed.
+        Thus no `model` is needed to be supplied when using this checkpoint.
+
+        Args:
+            file_path: The path to the .h5 file to load model from. This is the
+                same path that is used for ``model.save(path)``.
+            preprocessor: A fitted preprocessor to be applied before inference.
+
+        Returns:
+            A :py:class:`TensorflowCheckpoint` converted from h5 format.
+
+        """
+        if not os.path.isfile(file_path) or not file_path.endswith(".h5"):
+            raise ValueError(
+                "Please supply a h5 file path to `TensorflowCheckpoint.from_h5()`."
+            )
+        tempdir = tempfile.mkdtemp()
+        filename = os.path.basename(file_path)
+        new_checkpoint_file = Path(tempdir, filename).as_posix()
+        shutil.copy(file_path, new_checkpoint_file)
+
+        checkpoint = cls.from_directory(tempdir)
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+        checkpoint.update_metadata({cls.MODEL_FILENAME_KEY: filename})
+        return checkpoint
+
+    @classmethod
+    def from_saved_model(
+        cls, dir_path: str, *, preprocessor: Optional["Preprocessor"] = None
+    ) -> "TensorflowCheckpoint":
+        """Create a :py:class:`~ray.train.Checkpoint` that stores a Keras
+        model from SavedModel format.
+
+        The checkpoint generated by this method contains all the information needed.
+        Thus no `model` is needed to be supplied when using this checkpoint.
+
+        Args:
+            dir_path: The directory containing the saved model. This is the same
+                directory as used by ``model.save(dir_path)``.
+            preprocessor: A fitted preprocessor to be applied before inference.
+
+        Returns:
+            A :py:class:`TensorflowCheckpoint` converted from SavedModel format.
+
+        """
+        if not os.path.isdir(dir_path):
+            raise ValueError(
+                "Please supply a directory to `TensorflowCheckpoint.from_saved_model`"
+            )
+        tempdir = tempfile.mkdtemp()
+        # TODO(ml-team): Replace this with copytree()
+        os.rmdir(tempdir)
+        shutil.copytree(dir_path, tempdir)
+
+        checkpoint = cls.from_directory(tempdir)
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+        # NOTE: The entire directory is the checkpoint.
+        checkpoint.update_metadata({cls.MODEL_FILENAME_KEY: "."})
+        return checkpoint
+
+    def get_model(
+        self,
+    ) -> tf.keras.Model:
+        """Retrieve the model stored in this checkpoint.
+
+        Returns:
+            The Tensorflow Keras model stored in the checkpoint.
+        """
+        metadata = self.get_metadata()
+        if self.MODEL_FILENAME_KEY not in metadata:
+            raise ValueError(
+                "`TensorflowCheckpoint` cannot retrieve the model if you override the "
+                "checkpoint metadata. Please use `Checkpoint.update_metadata` instead."
+            )
+        model_filename = metadata[self.MODEL_FILENAME_KEY]
+        with self.as_directory() as checkpoint_dir:
+            model_path = Path(checkpoint_dir, model_filename).as_posix()
+            return keras.models.load_model(model_path)

python/ray/train/tensorflow/tensorflow_trainer.py

@@ -0,0 +1,191 @@
+from typing import Any, Callable, Dict, Optional, Union
+
+from ray.train import Checkpoint, DataConfig, RunConfig, ScalingConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tensorflow.config import TensorflowConfig
+from ray.train.trainer import GenDataset
+from ray.util import PublicAPI
+
+
+@PublicAPI(stability="beta")
+class TensorflowTrainer(DataParallelTrainer):
+    """A Trainer for data parallel Tensorflow training.
+
+    This Trainer runs the function ``train_loop_per_worker`` on multiple Ray
+    Actors. These actors already have the necessary TensorFlow process group already
+    configured for distributed TensorFlow training.
+
+    The ``train_loop_per_worker`` function is expected to take in either 0 or 1
+    arguments:
+
+    .. testcode::
+
+        def train_loop_per_worker():
+            ...
+
+    .. testcode::
+
+        def train_loop_per_worker(config: Dict):
+            ...
+
+    If ``train_loop_per_worker`` accepts an argument, then
+    ``train_loop_config`` will be passed in as the argument. This is useful if you
+    want to tune the values in ``train_loop_config`` as hyperparameters.
+
+    If the ``datasets`` dict contains a training dataset (denoted by
+    the "train" key), then it will be split into multiple dataset
+    shards that can then be accessed by ``ray.train.get_dataset_shard("train")`` inside
+    ``train_loop_per_worker``. All the other datasets will not be split and
+    ``ray.train.get_dataset_shard(...)`` will return the entire Dataset.
+
+    Inside the ``train_loop_per_worker`` function, you can use any of the
+    :ref:`Ray Train loop methods <train-loop-api>`.
+
+    .. warning::
+        Ray will not automatically set any environment variables or configuration
+        related to local parallelism / threading
+        :ref:`aside from "OMP_NUM_THREADS" <omp-num-thread-note>`.
+        If you desire greater control over TensorFlow threading, use
+        the ``tf.config.threading`` module (eg.
+        ``tf.config.threading.set_inter_op_parallelism_threads(num_cpus)``)
+        at the beginning of your ``train_loop_per_worker`` function.
+
+
+    .. testcode::
+
+        from ray import train
+
+        def train_loop_per_worker():
+            # Report intermediate results for callbacks or logging and
+            # checkpoint data.
+            train.report(...)
+
+            # Returns dict of last saved checkpoint.
+            train.get_checkpoint()
+
+            # Returns the Dataset shard for the given key.
+            train.get_dataset_shard("my_dataset")
+
+            # Returns the total number of workers executing training.
+            train.get_context().get_world_size()
+
+            # Returns the rank of this worker.
+            train.get_context().get_world_rank()
+
+            # Returns the rank of the worker on the current node.
+            train.get_context().get_local_rank()
+
+    Any returns from the ``train_loop_per_worker`` will be discarded and not
+    used or persisted anywhere.
+
+    Example:
+
+    .. testcode::
+
+        import os
+        import tempfile
+        import tensorflow as tf
+        import numpy as np
+
+        import ray
+        from ray import train
+        from ray.train import Checkpoint, ScalingConfig
+        from ray.train.tensorflow import TensorflowTrainer
+
+        def build_model():
+            # toy neural network : 1-layer
+            return tf.keras.Sequential(
+                [tf.keras.layers.Dense(
+                    1, activation="linear", input_shape=(1,))]
+            )
+
+        def train_loop_per_worker(config):
+            dataset_shard = train.get_dataset_shard("train")
+            strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+            with strategy.scope():
+                model = build_model()
+                model.compile(
+                    optimizer="Adam", loss="mean_squared_error", metrics=["mse"])
+
+            tf_dataset = dataset_shard.to_tf(
+                feature_columns="x",
+                label_columns="y",
+                batch_size=1
+            )
+            for epoch in range(config["num_epochs"]):
+                model.fit(tf_dataset)
+
+                # Create checkpoint.
+                checkpoint_dir = tempfile.mkdtemp()
+                model.save_weights(
+                    os.path.join(checkpoint_dir, "my_checkpoint")
+                )
+                checkpoint = Checkpoint.from_directory(checkpoint_dir)
+
+                train.report(
+                    {},
+                    checkpoint=checkpoint,
+                )
+
+        train_dataset = ray.data.from_items([{"x": np.array([x], dtype=np.float32), "y": x + 1} for x in range(32)])
+        trainer = TensorflowTrainer(
+            train_loop_per_worker=train_loop_per_worker,
+            scaling_config=ScalingConfig(num_workers=3, use_gpu=True),
+            datasets={"train": train_dataset},
+            train_loop_config={"num_epochs": 2},
+        )
+        result = trainer.fit()
+
+    .. testoutput::
+        :options:+ELLIPSIS
+        :hide:
+
+        ...
+
+    Args:
+        train_loop_per_worker: The training function to execute.
+            This can either take in no arguments or a ``config`` dict.
+        train_loop_config: Configurations to pass into
+            ``train_loop_per_worker`` if it accepts an argument.
+        tensorflow_config: Configuration for setting up the TensorFlow backend.
+            If set to None, use the default configuration. This replaces the
+            ``backend_config`` arg of ``DataParallelTrainer``.
+        scaling_config: Configuration for how to scale data parallel training.
+        dataset_config: Configuration for dataset ingest.
+        run_config: Configuration for the execution of the training run.
+        datasets: Any Datasets to use for training. Use
+            the key "train" to denote which dataset is the training
+            dataset.
+        metadata: Dict that should be made available via
+            `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
+            for checkpoints saved from this Trainer. Must be JSON-serializable.
+        resume_from_checkpoint: A checkpoint to resume training from.
+    """
+
+    def __init__(
+        self,
+        train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]],
+        *,
+        train_loop_config: Optional[Dict] = None,
+        tensorflow_config: Optional[TensorflowConfig] = None,
+        scaling_config: Optional[ScalingConfig] = None,
+        dataset_config: Optional[DataConfig] = None,
+        run_config: Optional[RunConfig] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        resume_from_checkpoint: Optional[Checkpoint] = None,
+    ):
+        if not tensorflow_config:
+            tensorflow_config = TensorflowConfig()
+
+        super(TensorflowTrainer, self).__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            train_loop_config=train_loop_config,
+            backend_config=tensorflow_config,
+            scaling_config=scaling_config,
+            dataset_config=dataset_config,
+            run_config=run_config,
+            datasets=datasets,
+            resume_from_checkpoint=resume_from_checkpoint,
+            metadata=metadata,
+        )

python/ray/train/tensorflow/train_loop_utils.py

@@ -0,0 +1,27 @@
+import tensorflow as tf
+
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="beta")
+def prepare_dataset_shard(tf_dataset_shard: tf.data.Dataset):
+    """A utility function that overrides default config for Tensorflow Dataset.
+
+    This should be used on a TensorFlow ``Dataset`` created by calling
+    ``iter_tf_batches()`` on a ``ray.data.Dataset`` returned by
+    ``ray.train.get_dataset_shard()`` since the dataset has already
+    been sharded across the workers.
+
+    Args:
+        tf_dataset_shard: A TensorFlow Dataset.
+
+    Returns:
+        A TensorFlow Dataset with:
+            - autosharding turned off
+            - prefetching turned on with autotune enabled
+    """
+    options = tf.data.Options()
+    options.experimental_distribute.auto_shard_policy = (
+        tf.data.experimental.AutoShardPolicy.OFF
+    )
+    return tf_dataset_shard.with_options(options).prefetch(tf.data.AUTOTUNE)

python/ray/train/tests/conftest.py

@@ -0,0 +1,143 @@
+import logging
+
+import boto3
+import pytest
+
+import ray
+from ray._common.test_utils import simulate_s3_bucket
+from ray.cluster_utils import Cluster
+
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+from ray.tests.conftest import (
+    propagate_logs,  # noqa
+    pytest_runtest_makereport,  # noqa
+)
+
+
+@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()
+
+
+@pytest.fixture
+def ray_start_1_cpu_1_gpu():
+    address_info = ray.init(num_cpus=1, num_gpus=1)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_4_cpus_2_gpus():
+    address_info = ray.init(num_cpus=4, num_gpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def shutdown_only():
+    yield None
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_2_node_2_gpu():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=4, num_gpus=2)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_2_node_2_neuron_cores():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=4, resources={"neuron_cores": 2})
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@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()
+
+
+@pytest.fixture
+def ray_4_node_4_cpu():
+    cluster = Cluster()
+    for _ in range(4):
+        cluster.add_node(num_cpus=4)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_2_node_4_gpu():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=2, num_gpus=4)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_2_node_2_cpu():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=2)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def mock_s3_bucket_uri():
+    from ray.air._internal.uri_utils import URI
+
+    port = 5002
+    region = "us-west-2"
+    with simulate_s3_bucket(port=port, region=region) as s3_uri:
+        s3 = boto3.client(
+            "s3", region_name=region, endpoint_url=f"http://localhost:{port}"
+        )
+        # Bucket name will be autogenerated/unique per test
+        bucket_name = URI(s3_uri).name
+        s3.create_bucket(
+            Bucket=bucket_name,
+            CreateBucketConfiguration={"LocationConstraint": region},
+        )
+        # Disable server HTTP request logging
+        logging.getLogger("werkzeug").setLevel(logging.WARNING)
+        yield s3_uri
+        logging.getLogger("werkzeug").setLevel(logging.INFO)

python/ray/train/tests/dummy_preprocessor.py

@@ -0,0 +1,29 @@
+import uuid
+
+from ray.data.preprocessor import Preprocessor
+
+
+class DummyPreprocessor(Preprocessor):
+    _is_fittable = False
+
+    def __init__(self, transform=None):
+        self.id = uuid.uuid4()
+
+        if transform is None:
+            self.transform = lambda b: b
+        else:
+            self.transform = transform
+
+    def transform_batch(self, batch):
+        self._batch_transformed = True
+        return self.transform(batch)
+
+    def _transform_pandas(self, df):
+        return df
+
+    @property
+    def has_preprocessed(self):
+        return hasattr(self, "_batch_transformed")
+
+    def __eq__(self, other_preprocessor):
+        return self.id == other_preprocessor.id

python/ray/train/tests/lightning_test_utils.py

@@ -0,0 +1,171 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchmetrics import Accuracy
+
+from ray import train
+from ray.train.lightning._lightning_utils import import_lightning
+
+pl = import_lightning()
+
+
+class LinearModule(pl.LightningModule):
+    def __init__(self, input_dim, output_dim, strategy="ddp", fail_epoch=-1) -> None:
+        super().__init__()
+        self.linear = nn.Linear(input_dim, output_dim)
+        self.loss = []
+        self.strategy = strategy
+        self.restored = train.get_checkpoint() is not None
+        self.fail_epoch = fail_epoch
+
+    def forward(self, input):
+        if isinstance(input, dict) and len(input) == 1:
+            input = list(input.values())[0]
+        return self.linear(input)
+
+    def training_step(self, batch):
+        if not self.restored and self.fail_epoch == self.current_epoch:
+            raise RuntimeError
+
+        output = self.forward(batch)
+        loss = torch.sum(output)
+        self.log("loss", loss)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        loss = self.forward(val_batch)
+        self.loss.append(loss)
+        return {"val_loss": loss}
+
+    def test_step(self, batch, batch_idx):
+        loss = self.forward(batch)
+        return {"test_loss": loss}
+
+    def on_validation_epoch_end(self) -> None:
+        avg_loss = torch.stack(self.loss).mean()
+        self.log("val_loss", avg_loss)
+        self.loss.clear()
+
+    def predict_step(self, batch, batch_idx):
+        return self.forward(batch)
+
+    def configure_optimizers(self):
+        if self.strategy == "fsdp":
+            # Feed FSDP wrapped model parameters to optimizer
+            return torch.optim.AdamW(self.trainer.model.parameters(), lr=0.1)
+        else:
+            return torch.optim.AdamW(self.parameters(), lr=0.1)
+
+
+class DoubleLinearModule(pl.LightningModule):
+    def __init__(self, input_dim_1, input_dim_2, output_dim) -> None:
+        super().__init__()
+        self.linear_1 = nn.Linear(input_dim_1, output_dim)
+        self.linear_2 = nn.Linear(input_dim_2, output_dim)
+        self.loss = []
+
+    def forward(self, batch):
+        input_1 = batch["input_1"]
+        input_2 = batch["input_2"]
+        return self.linear_1(input_1) + self.linear_2(input_2)
+
+    def training_step(self, batch):
+        output = self.forward(batch)
+        loss = torch.sum(output)
+        self.log("loss", loss)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        loss = self.forward(val_batch)
+        self.loss.append(loss)
+        return {"val_loss": loss}
+
+    def on_validation_epoch_end(self) -> None:
+        print("Validation Epoch:", self.current_epoch)
+        avg_loss = torch.stack(self.loss).mean()
+        self.log("val_loss", avg_loss)
+        self.loss.clear()
+
+    def predict_step(self, batch, batch_idx):
+        return self.forward(batch)
+
+    def configure_optimizers(self):
+        return torch.optim.AdamW(self.parameters(), lr=0.1)
+
+
+class DummyDataModule(pl.LightningDataModule):
+    def __init__(self, batch_size: int = 8, dataset_size: int = 256) -> None:
+        super().__init__()
+        self.batch_size = batch_size
+        self.train_data = torch.randn(dataset_size, 32)
+        self.val_data = torch.randn(dataset_size, 32)
+        self.test_data = torch.randn(dataset_size, 32)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_data, batch_size=self.batch_size)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_data, batch_size=self.batch_size)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_data, batch_size=self.batch_size)
+
+
+class LightningMNISTClassifier(pl.LightningModule):
+    def __init__(self, lr: float, layer_1: int, layer_2: int):
+        super(LightningMNISTClassifier, self).__init__()
+
+        self.lr = lr
+        # mnist images are (1, 28, 28) (channels, width, height)
+        self.layer_1 = torch.nn.Linear(28 * 28, layer_1)
+        self.layer_2 = torch.nn.Linear(layer_1, layer_2)
+        self.layer_3 = torch.nn.Linear(layer_2, 10)
+        self.accuracy = Accuracy(task="multiclass", num_classes=10, top_k=1)
+        self.val_acc_list = []
+        self.val_loss_list = []
+
+    def forward(self, x):
+        batch_size, channels, width, height = x.size()
+        x = x.view(batch_size, -1)
+        x = self.layer_1(x)
+        x = torch.relu(x)
+        x = self.layer_2(x)
+        x = torch.relu(x)
+        x = self.layer_3(x)
+        x = torch.log_softmax(x, dim=1)
+        return x
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.lr)
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        logits = self.forward(x)
+        loss = F.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        self.log("ptl/train_loss", loss)
+        self.log("ptl/train_accuracy", acc)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        logits = self.forward(x)
+        loss = F.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        self.val_acc_list.append(acc)
+        self.val_loss_list.append(loss)
+        return {"val_loss": loss, "val_accuracy": acc}
+
+    def on_validation_epoch_end(self):
+        avg_loss = torch.stack(self.val_loss_list).mean()
+        avg_acc = torch.stack(self.val_acc_list).mean()
+        self.log("ptl/val_loss", avg_loss)
+        self.log("ptl/val_accuracy", avg_acc)
+        self.val_acc_list.clear()
+        self.val_loss_list.clear()
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=None):
+        x = batch
+        logits = self.forward(x)
+        return torch.argmax(logits, dim=-1)

python/ray/train/tests/test_api_migrations.py

@@ -0,0 +1,158 @@
+import sys
+import warnings
+
+import pytest
+
+import ray.train
+import ray.tune
+from ray.train.constants import ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.util.annotations import RayDeprecationWarning
+
+
+@pytest.fixture(autouse=True)
+def enable_v2_migration_deprecation_messages(monkeypatch):
+    monkeypatch.setenv(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR, "1")
+    yield
+    monkeypatch.delenv(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR)
+
+
+def test_trainer_restore():
+    with pytest.warns(RayDeprecationWarning, match="restore"):
+        try:
+            DataParallelTrainer.restore("dummy")
+        except Exception:
+            pass
+
+    with pytest.warns(RayDeprecationWarning, match="can_restore"):
+        try:
+            DataParallelTrainer.can_restore("dummy")
+        except Exception:
+            pass
+
+
+def test_trainer_valid_configs(ray_start_4_cpus, tmp_path):
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        DataParallelTrainer(
+            lambda _: None,
+            scaling_config=ray.train.ScalingConfig(num_workers=1),
+            run_config=ray.train.RunConfig(
+                storage_path=tmp_path,
+                failure_config=ray.train.FailureConfig(max_failures=1),
+            ),
+        ).fit()
+
+    for warning in w:
+        assert not (
+            warning.category == RayDeprecationWarning
+            and "`RunConfig` class should be imported from `ray.tune`"
+            in str(warning.message)
+        )
+
+
+def test_trainer_deprecated_configs():
+    with pytest.warns(RayDeprecationWarning, match="metadata"):
+        DataParallelTrainer(
+            lambda _: None,
+            metadata={"dummy": "dummy"},
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="resume_from_checkpoint"):
+        DataParallelTrainer(
+            lambda _: None,
+            resume_from_checkpoint=ray.train.Checkpoint.from_directory("dummy"),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="fail_fast"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(
+                failure_config=ray.train.FailureConfig(fail_fast=True)
+            ),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="trainer_resources"):
+        DataParallelTrainer(
+            lambda _: None,
+            scaling_config=ray.train.ScalingConfig(trainer_resources={"CPU": 1}),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="verbose"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(verbose=True),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="log_to_file"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(log_to_file=True),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="stop"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(stop={"training_iteration": 1}),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="callbacks"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(callbacks=[ray.tune.Callback()]),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="progress_reporter"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(
+                progress_reporter=ray.tune.ProgressReporter()
+            ),
+        )
+
+    with pytest.warns(RayDeprecationWarning, match="sync_config"):
+        DataParallelTrainer(
+            lambda _: None,
+            run_config=ray.train.RunConfig(
+                sync_config=ray.train.SyncConfig(sync_artifacts=True)
+            ),
+        )
+
+
+def test_train_context_deprecations(ray_start_4_cpus, tmp_path):
+    def train_fn_per_worker(config):
+        with pytest.warns(RayDeprecationWarning, match="get_trial_dir"):
+            ray.train.get_context().get_trial_dir()
+
+        with pytest.warns(RayDeprecationWarning, match="get_trial_id"):
+            ray.train.get_context().get_trial_id()
+
+        with pytest.warns(RayDeprecationWarning, match="get_trial_name"):
+            ray.train.get_context().get_trial_name()
+
+        with pytest.warns(RayDeprecationWarning, match="get_trial_resources"):
+            ray.train.get_context().get_trial_resources()
+
+    trainer = DataParallelTrainer(
+        train_fn_per_worker,
+        scaling_config=ray.train.ScalingConfig(num_workers=1),
+        run_config=ray.train.RunConfig(storage_path=tmp_path),
+    )
+    trainer.fit()
+
+
+def test_v2_enabled_error(monkeypatch):
+    """Running a V1 Trainer with V2 enabled should raise an error."""
+    from ray.train.v2._internal.constants import V2_ENABLED_ENV_VAR
+
+    monkeypatch.setenv(V2_ENABLED_ENV_VAR, "1")
+
+    with pytest.raises(DeprecationWarning, match="Detected use of a deprecated"):
+        DataParallelTrainer(
+            lambda _: None,
+            scaling_config=ray.train.ScalingConfig(num_workers=1),
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_backend.py

@@ -0,0 +1,660 @@
+import math
+import os
+import sys
+import tempfile
+import time
+from typing import Set
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray import train
+from ray._private.accelerators.neuron import NEURON_RT_VISIBLE_CORES_ENV_VAR
+from ray.air._internal.util import StartTraceback
+
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+from ray.tests.conftest import pytest_runtest_makereport  # noqa
+from ray.train import DataConfig
+from ray.train._internal.backend_executor import (
+    BackendExecutor,
+    InactiveWorkerGroupError,
+    TrainBackendError,
+    TrainingWorkerError,
+)
+from ray.train._internal.storage import StorageContext
+from ray.train._internal.worker_group import WorkerGroup, WorkerMetadata
+from ray.train.backend import Backend, BackendConfig
+from ray.train.constants import (
+    ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
+    ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
+    JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
+    TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
+    TRAIN_ENABLE_WORKER_SPREAD_ENV,
+)
+from ray.train.torch import TorchConfig
+from ray.train.v2.jax.config import JaxConfig
+from ray.util.placement_group import get_current_placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from ray.util.state import list_actors
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    ray.shutdown()
+
+
+def _start_training(backend_executor: BackendExecutor, fn):
+    storage = StorageContext(
+        storage_path=tempfile.mkdtemp(),
+        experiment_dir_name="exp_name",
+        trial_dir_name="trial_name",
+    )
+    backend_executor.start_training(
+        train_func=fn,
+        datasets={},
+        metadata={},
+        data_config=DataConfig(),
+        storage=storage,
+    )
+
+
+def gen_execute_special(special_f):
+    def execute_async_special(self, f):
+        """Runs f on worker 0, special_f on other workers."""
+        futures = [
+            self.workers[0]
+            .actor._RayTrainWorker__execute.options(name=f.__name__)
+            .remote(f)
+        ]
+        for worker in self.workers[1:]:
+            futures.append(
+                worker.actor._RayTrainWorker__execute.options(
+                    name=special_f.__name__
+                ).remote(special_f)
+            )
+        return futures
+
+    return execute_async_special
+
+
+class TestConfig(BackendConfig):
+    @property
+    def backend_cls(self):
+        return TestBackend
+
+
+class TestBackend(Backend):
+    def on_start(self, worker_group: WorkerGroup, backend_config: TestConfig):
+        pass
+
+    def on_shutdown(self, worker_group: WorkerGroup, backend_config: TestConfig):
+        pass
+
+
+original_add_workers = WorkerGroup.add_workers
+
+
+def mock_add_workers(self, num_workers):
+    original_add_workers(self, num_workers)
+    for i, worker in enumerate(self.workers):
+        metadata = WorkerMetadata(
+            node_id=str(i % 2),
+            node_ip=str(i % 2),
+            hostname=0,
+            resource_ids={"GPU": ["0"]},
+            pid=0,
+        )
+        worker.metadata = metadata
+
+
+def mock_add_workers_to_nodes_with_same_ip(self, num_workers):
+    original_add_workers(self, num_workers)
+    for i, worker in enumerate(self.workers):
+        metadata = WorkerMetadata(
+            node_id=str(i % 2),
+            node_ip=0,
+            hostname=0,
+            resource_ids={"GPU": ["0"]},
+            pid=0,
+        )
+        worker.metadata = metadata
+
+
+def test_start(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    with pytest.raises(InactiveWorkerGroupError):
+        _start_training(e, lambda: 1)
+    e.start()
+    assert len(e.worker_group) == 2
+
+
+def test_initialization_hook(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+
+    def init_hook():
+        import os
+
+        os.environ["TEST"] = "1"
+
+    e.start(initialization_hook=init_hook)
+
+    def check():
+        import os
+
+        return os.getenv("TEST", "0")
+
+    _start_training(e, check)
+    assert e.finish_training() == ["1", "1"]
+
+
+def test_shutdown(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+    assert len(e.worker_group) == 2
+    e.shutdown()
+    with pytest.raises(InactiveWorkerGroupError):
+        _start_training(e, lambda: 1)
+
+
+def test_train(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    _start_training(e, lambda: 1)
+    assert e.finish_training() == [1, 1]
+
+
+def test_local_ranks(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    def train_func():
+        return train.get_context().get_local_rank()
+
+    _start_training(e, train_func)
+    assert set(e.finish_training()) == {0, 1}
+
+
+def test_local_ranks_with_same_ip_nodes(ray_2_node_2_cpu):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=4)
+    e.start()
+
+    def train_func():
+        return train.get_context().get_local_rank()
+
+    _start_training(e, train_func)
+    assert list(e.finish_training()) == [0, 1, 0, 1]
+
+
+def test_local_world_size(ray_2_node_2_cpu):
+    config = TestConfig()
+    with patch.object(WorkerGroup, "add_workers", mock_add_workers):
+        e = BackendExecutor(config, num_workers=3)
+        e.start()
+
+        def train_func():
+            return train.get_context().get_local_world_size()
+
+        _start_training(e, train_func)
+        assert list(e.finish_training()) == [2, 2, 1]
+
+
+def test_local_world_size_with_same_ip_nodes(ray_2_node_2_cpu):
+    config = TestConfig()
+    with patch.object(
+        WorkerGroup, "add_workers", mock_add_workers_to_nodes_with_same_ip
+    ):
+        e = BackendExecutor(config, num_workers=3)
+        e.start()
+
+        def train_func():
+            return train.get_context().get_local_world_size()
+
+        _start_training(e, train_func)
+        assert list(e.finish_training()) == [2, 2, 1]
+
+
+def test_node_ranks(ray_2_node_2_cpu):
+    config = TestConfig()
+    with patch.object(WorkerGroup, "add_workers", mock_add_workers):
+        e = BackendExecutor(config, num_workers=3)
+        e.start()
+
+        def train_func():
+            return train.get_context().get_node_rank()
+
+        _start_training(e, train_func)
+        assert list(e.finish_training()) == [0, 0, 1]
+
+
+def test_node_ranks_with_same_ip_nodes(ray_2_node_2_cpu):
+    config = TestConfig()
+    with patch.object(
+        WorkerGroup, "add_workers", mock_add_workers_to_nodes_with_same_ip
+    ):
+        e = BackendExecutor(config, num_workers=3)
+        e.start()
+
+        def train_func():
+            return train.get_context().get_node_rank()
+
+        _start_training(e, train_func)
+        assert list(e.finish_training()) == [0, 0, 1]
+
+
+def test_train_failure(ray_start_2_cpus):
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    with pytest.raises(StartTraceback) as exc:
+        e.get_next_results()
+    assert isinstance(exc.value.__cause__, TrainBackendError)
+
+    with pytest.raises(StartTraceback) as exc:
+        e.pause_reporting()
+    assert isinstance(exc.value.__cause__, TrainBackendError)
+
+    with pytest.raises(StartTraceback) as exc:
+        e.finish_training()
+    assert isinstance(exc.value.__cause__, TrainBackendError)
+
+    _start_training(e, lambda: 1)
+
+    with pytest.raises(StartTraceback) as exc:
+        _start_training(e, lambda: 2)
+    assert isinstance(exc.value.__cause__, TrainBackendError)
+
+    assert e.finish_training() == [1, 1]
+
+
+def test_single_worker_user_failure(ray_start_2_cpus):
+    """Tests if training fails immediately if one worker raises an Exception
+    while executing the user training code."""
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    def single_worker_user_failure():
+        if train.get_context().get_world_rank() == 0:
+            raise RuntimeError
+        else:
+            time.sleep(1000000)
+
+    _start_training(e, single_worker_user_failure)
+
+    with pytest.raises(StartTraceback) as exc:
+        e.get_next_results()
+    assert isinstance(exc.value.__cause__, RuntimeError)
+
+
+def test_single_worker_actor_failure(ray_start_2_cpus):
+    """Tests is training fails immediately if one worker actor dies."""
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=2)
+    e.start()
+
+    def single_worker_actor_failure():
+        if train.get_context().get_world_rank() == 0:
+            # Simulate actor failure
+            os._exit(1)
+        else:
+            time.sleep(1000)
+
+    _start_training(e, single_worker_actor_failure)
+
+    with pytest.raises(TrainingWorkerError):
+        e.get_next_results()
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="tensorflow is not supported in python 3.12+"
+)
+def test_tensorflow_start(ray_start_2_cpus):
+    from ray.train.tensorflow import TensorflowConfig
+
+    num_workers = 2
+    tensorflow_config = TensorflowConfig()
+    e = BackendExecutor(tensorflow_config, num_workers=num_workers)
+    e.start()
+
+    def get_tf_config():
+        import json
+        import os
+
+        return json.loads(os.environ["TF_CONFIG"])
+
+    _start_training(e, get_tf_config)
+    results = e.finish_training()
+    assert len(results) == num_workers
+
+    workers = [result["cluster"]["worker"] for result in results]
+    assert all(worker == workers[0] for worker in workers)
+
+    indexes = [result["task"]["index"] for result in results]
+    assert len(set(indexes)) == num_workers
+
+
+@pytest.mark.parametrize("init_method", ["env", "tcp"])
+def test_torch_start_shutdown(ray_start_2_cpus, init_method):
+    torch_config = TorchConfig(backend="gloo", init_method=init_method)
+    e = BackendExecutor(torch_config, num_workers=2)
+    e.start()
+
+    def check_process_group():
+        import torch
+
+        return (
+            torch.distributed.is_initialized()
+            and torch.distributed.get_world_size() == 2
+        )
+
+    _start_training(e, check_process_group)
+    assert all(e.finish_training())
+
+    e._backend.on_shutdown(e.worker_group, e._backend_config)
+
+    _start_training(e, check_process_group)
+    assert not any(e.finish_training())
+
+
+@pytest.mark.parametrize(
+    "init_method, timeout_s", [("env", 5), ("tcp", 5), ("env", 0), ("tcp", 0)]
+)
+def test_torch_process_group_shutdown_timeout(
+    ray_start_2_cpus, monkeypatch, init_method, timeout_s
+):
+    monkeypatch.setenv(TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, timeout_s)
+    torch_config = TorchConfig(backend="gloo", init_method=init_method)
+    e = BackendExecutor(torch_config, num_workers=2)
+    e.start()
+
+    _start_training(e, lambda: 1)
+    assert e.finish_training() == [1, 1]
+
+    # Verify that we do not raise an exception even if we time out
+    e._backend.on_shutdown(e.worker_group, e._backend_config)
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0]]),
+        (2, [[0, 1]] * 2),
+        (3, [[0]] + [[0, 1]] * 2),
+        (4, [[0, 1]] * 4),
+    ],
+)
+def test_cuda_visible_devices(ray_2_node_2_gpu, worker_results):
+    config = TestConfig()
+
+    def get_resources():
+        cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        # Sort the cuda visible devices to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(cuda_visible_devices.split(","))
+        ]
+        return sorted_devices
+
+    num_workers, expected_results = worker_results
+
+    os.environ[ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV] = "1"
+    e = BackendExecutor(
+        config, num_workers=num_workers, resources_per_worker={"GPU": 1}
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0]]),
+        (
+            2,
+            [[0]] * 2,
+        ),
+        (3, [[0, 1]] * 3),
+        (4, [[0, 1]] * 4),
+        (5, [[0]] + [[0, 1]] * 4),
+        (6, [[0]] * 2 + [[0, 1]] * 4),
+        (7, [[0, 1]] * 7),
+        (8, [[0, 1]] * 8),
+    ],
+)
+def test_cuda_visible_devices_fractional(ray_2_node_2_gpu, worker_results):
+    config = TestConfig()
+
+    if worker_results[0] != len(worker_results[1]):
+        raise ValueError(
+            "Invalid test parameter. Length of expected result should "
+            "match number of workers."
+        )
+
+    def get_resources():
+        cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        # Sort the cuda visible devices to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(cuda_visible_devices.split(","))
+        ]
+        return sorted_devices
+
+    num_workers, expected_results = worker_results
+
+    os.environ[ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV] = "1"
+    e = BackendExecutor(
+        config, num_workers=num_workers, resources_per_worker={"GPU": 0.5}
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0, 1]]),
+        (2, [[0, 1, 2, 3]] * 2),
+        (3, [[0, 1]] + [[0, 1, 2, 3]] * 2),
+        (4, [[0, 1, 2, 3]] * 4),
+    ],
+)
+def test_cuda_visible_devices_multiple(ray_2_node_4_gpu, worker_results):
+    config = TestConfig()
+
+    def get_resources():
+        cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        # Sort the cuda visible devices to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(cuda_visible_devices.split(","))
+        ]
+        return sorted_devices
+
+    if worker_results[0] != len(worker_results[1]):
+        raise ValueError(
+            "Invalid test parameter. Length of expected result should "
+            "match number of workers."
+        )
+
+    num_workers, expected_results = worker_results
+
+    os.environ[ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV] = "1"
+    e = BackendExecutor(
+        config, num_workers=num_workers, resources_per_worker={"GPU": 2}
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0]]),
+        (2, [[0, 1]] * 2),
+        (3, [[0]] + [[0, 1]] * 2),
+        (4, [[0, 1]] * 4),
+    ],
+)
+def test_neuron_core_accelerator_ids(ray_2_node_2_neuron_cores, worker_results):
+    config = TestConfig()
+
+    def get_resources():
+        neuron_resource_ids = os.environ[NEURON_RT_VISIBLE_CORES_ENV_VAR]
+        # Sort the runtime ids to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(neuron_resource_ids.split(","))
+        ]
+        return sorted_devices
+
+    num_workers, expected_results = worker_results
+    # sharing enabled by default
+    os.environ.pop(ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV, None)
+    e = BackendExecutor(
+        config,
+        num_workers=num_workers,
+        resources_per_worker={"neuron_cores": 1},
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+@pytest.mark.parametrize(
+    "worker_results",
+    [
+        (1, [[0]]),
+        (2, [[0]] + [[1]]),
+        (3, [[0]] * 2 + [[1]]),
+        (4, [[0]] * 2 + [[1]] * 2),
+    ],
+)
+def test_neuron_core_accelerator_ids_sharing_disabled(
+    ray_2_node_2_neuron_cores, worker_results
+):
+    config = TestConfig()
+
+    def get_resources():
+        neuron_resource_ids = os.environ[NEURON_RT_VISIBLE_CORES_ENV_VAR]
+        # Sort the runtime ids to have exact match with expected result.
+        sorted_devices = [
+            int(device) for device in sorted(neuron_resource_ids.split(","))
+        ]
+        return sorted_devices
+
+    num_workers, expected_results = worker_results
+
+    os.environ[ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV] = "0"
+    e = BackendExecutor(
+        config,
+        num_workers=num_workers,
+        resources_per_worker={"neuron_cores": 1},
+    )
+    e.start()
+    _start_training(e, get_resources)
+    results = e.finish_training()
+    results.sort()
+    assert results == expected_results
+
+
+def get_node_id_set() -> Set[str]:
+    return {a.node_id for a in list_actors()}
+
+
+@pytest.mark.parametrize("num_workers", [3, 4, 5])
+def test_placement_group_pack(ray_4_node_4_cpu, num_workers):
+    """Tests that workers are packed on nodes."""
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=num_workers)
+    e.start()
+    node_id_set = get_node_id_set()
+    assert len(node_id_set) == math.ceil(num_workers / 4)
+
+
+@pytest.mark.parametrize("num_workers", [3, 4, 5])
+def test_placement_group_spread(ray_4_node_4_cpu, num_workers):
+    """Tests that workers are spread across nodes."""
+    os.environ[TRAIN_ENABLE_WORKER_SPREAD_ENV] = "1"
+    config = TestConfig()
+    e = BackendExecutor(config, num_workers=num_workers)
+    e.start()
+    node_id_set = get_node_id_set()
+    assert len(node_id_set) == min(num_workers, 4)
+
+
+@pytest.mark.parametrize("placement_group_capture_child_tasks", [True, False])
+def test_placement_group_parent(ray_4_node_4_cpu, placement_group_capture_child_tasks):
+    """Tests that parent placement group will be used."""
+    num_workers = 2
+    bundle = {"CPU": 1}
+    bundles = [bundle.copy() for _ in range(num_workers + 1)]
+    placement_group = ray.util.placement_group(bundles)
+
+    def train_func():
+        return get_current_placement_group().id
+
+    @ray.remote
+    def test():
+        config = TestConfig()
+        e = BackendExecutor(config, num_workers=2)
+        e.start()
+        _start_training(e, train_func)
+        return e.finish_training()
+
+    results_future = test.options(
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=placement_group,
+            placement_group_capture_child_tasks=placement_group_capture_child_tasks,
+        ),
+    ).remote()
+    results = ray.get(results_future)
+    for worker_result in results:
+        if placement_group_capture_child_tasks:
+            assert worker_result == placement_group.id
+        else:
+            assert worker_result != placement_group.id
+
+
+@pytest.mark.parametrize("timeout_s", [5, 0])
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Current jax version is not supported in python 3.12+",
+)
+def test_jax_distributed_shutdown_timeout(ray_start_2_cpus, monkeypatch, timeout_s):
+    """Test that JAX distributed shutdown respects the timeout env var."""
+    monkeypatch.setenv(JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S, str(timeout_s))
+    jax_config = JaxConfig(use_tpu=True)
+    e = BackendExecutor(jax_config, num_workers=2)
+    e.start()
+
+    _start_training(e, lambda: 1)
+    assert e.finish_training() == [1, 1]
+
+    # Verify that we do not raise an exception even if we time out
+    e._backend.on_shutdown(e.worker_group, e._backend_config)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/train/tests/test_base_trainer.py

@@ -0,0 +1,197 @@
+import logging
+import tempfile
+
+import numpy as np
+import pytest
+
+import ray
+from ray import train, tune
+from ray.data.context import DataContext
+from ray.train import Checkpoint, ScalingConfig
+from ray.train._internal.session import get_session
+from ray.train.base_trainer import format_datasets_for_repr
+from ray.train.trainer import BaseTrainer
+from ray.util.placement_group import get_current_placement_group
+
+logger = logging.getLogger(__name__)
+
+
+class DummyTrainer(BaseTrainer):
+    _scaling_config_allowed_keys = BaseTrainer._scaling_config_allowed_keys + [
+        "num_workers",
+        "use_gpu",
+        "resources_per_worker",
+        "placement_strategy",
+    ]
+
+    def __init__(self, train_loop, custom_arg=None, **kwargs):
+        self.custom_arg = custom_arg
+        self.train_loop = train_loop
+        super().__init__(**kwargs)
+
+    def training_loop(self) -> None:
+        self.train_loop(self)
+
+
+def test_trainer_fit(ray_start_4_cpus):
+    def training_loop(self):
+        train.report(dict(my_metric=1))
+
+    trainer = DummyTrainer(train_loop=training_loop)
+    result = trainer.fit()
+    assert result.metrics["my_metric"] == 1
+
+
+def test_validate_datasets(ray_start_4_cpus):
+    with pytest.raises(ValueError) as e:
+        DummyTrainer(train_loop=None, datasets=1)
+    assert "`datasets` should be a dict mapping" in str(e.value)
+
+    with pytest.raises(ValueError) as e:
+        DummyTrainer(train_loop=None, datasets={"train": 1})
+    assert "The Dataset under train key is not a `ray.data.Dataset`"
+
+
+def test_resources(ray_start_4_cpus):
+    def check_cpus(self):
+        assert ray.available_resources()["CPU"] == 2
+
+    assert ray.available_resources()["CPU"] == 4
+    trainer = DummyTrainer(
+        check_cpus,
+        scaling_config=ScalingConfig(
+            trainer_resources={"CPU": 2}, resources_per_worker={}
+        ),
+    )
+    trainer.fit()
+
+
+def test_arg_override(ray_start_4_cpus):
+    def check_override(self):
+        assert self.scaling_config.num_workers == 1
+        # Should do deep update.
+        assert not self.custom_arg["outer"]["inner"]
+        assert self.custom_arg["outer"]["fixed"] == 1
+
+        pg = get_current_placement_group()
+        assert len(pg.bundle_specs) == 1  # Merged trainer and worker bundle
+
+    scale_config = ScalingConfig(num_workers=4)
+    trainer = DummyTrainer(
+        check_override,
+        custom_arg={"outer": {"inner": True, "fixed": 1}},
+        scaling_config=scale_config,
+    )
+
+    new_config = {
+        "custom_arg": {"outer": {"inner": False}},
+        "scaling_config": ScalingConfig(num_workers=1),
+    }
+
+    tune.run(trainer.as_trainable(), config=new_config)
+
+
+def test_reserved_cpu_warnings_no_cpu_usage(ray_start_1_cpu_1_gpu):
+    """Ensure there is no divide by zero error if trial requires no CPUs."""
+
+    def train_loop(config):
+        pass
+
+    trainer = DummyTrainer(
+        train_loop,
+        scaling_config=ScalingConfig(
+            num_workers=1, use_gpu=True, trainer_resources={"CPU": 0}
+        ),
+        datasets={"train": ray.data.range(10)},
+    )
+    trainer.fit()
+
+
+def test_setup(ray_start_4_cpus):
+    def check_setup(self):
+        assert self._has_setup
+
+    class DummyTrainerWithSetup(DummyTrainer):
+        def setup(self):
+            self._has_setup = True
+
+    trainer = DummyTrainerWithSetup(check_setup)
+    trainer.fit()
+
+
+def test_repr(ray_start_4_cpus):
+    def training_loop(self):
+        pass
+
+    trainer = DummyTrainer(
+        training_loop,
+        datasets={
+            "train": ray.data.from_items([1, 2, 3]),
+        },
+    )
+
+    representation = repr(trainer)
+
+    assert "DummyTrainer" in representation
+
+
+def test_metadata_propagation(ray_start_4_cpus):
+    class MyTrainer(BaseTrainer):
+        def training_loop(self):
+            assert get_session().metadata == {"a": 1, "b": 1}
+            with tempfile.TemporaryDirectory() as path:
+                checkpoint = Checkpoint.from_directory(path)
+                checkpoint.set_metadata({"b": 2, "c": 3})
+                train.report(dict(my_metric=1), checkpoint=checkpoint)
+
+    trainer = MyTrainer(metadata={"a": 1, "b": 1})
+    result = trainer.fit()
+    meta_out = result.checkpoint.get_metadata()
+    assert meta_out == {"a": 1, "b": 2, "c": 3}, meta_out
+
+
+def test_data_context_propagation(ray_start_4_cpus):
+    ctx = DataContext.get_current()
+    # Fake DataContext attribute to propagate to worker.
+    ctx.foo = "bar"
+
+    def training_loop(self):
+        # Dummy train loop that checks that changes in the driver's
+        # DataContext are propagated to the worker.
+        ctx_worker = DataContext.get_current()
+        assert ctx_worker.foo == "bar"
+
+    trainer = DummyTrainer(
+        train_loop=training_loop,
+        datasets={"train": ray.data.range(10)},
+    )
+    trainer.fit()
+
+
+def test_large_params(ray_start_4_cpus):
+    """Tests that large params are not serialized with the trainer actor
+    and are instead put into the object store separately."""
+    huge_array = np.zeros(shape=int(1e8))
+
+    def training_loop(self):
+        _ = huge_array
+
+    trainer = DummyTrainer(training_loop)
+    trainer.fit()
+
+
+def test_format_datasets_for_repr(ray_start_4_cpus):
+    datasets = {"train": ray.data.range(1), "test": ray.data.range(1)}
+
+    actual_repr = format_datasets_for_repr(datasets)
+
+    assert actual_repr == (
+        "{'train': Dataset(num_rows=1, schema={id: int64}), "
+        "'test': Dataset(num_rows=1, schema={id: int64})}"
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(sys.argv[1:] + ["-v", "-x", __file__]))

python/ray/train/tests/test_base_worker_group.py

@@ -0,0 +1,46 @@
+"""Tests for BaseWorkerGroup implementation and usage."""
+
+import pytest
+
+from ray.train._internal.base_worker_group import BaseWorkerGroup
+from ray.train._internal.worker_group import WorkerGroup as V1WorkerGroup
+from ray.train.v2._internal.execution.worker_group.worker_group import (
+    WorkerGroup as V2WorkerGroup,
+)
+
+
+def test_interface_abstract_methods():
+    """Test that BaseWorkerGroup enforces its abstract methods."""
+    # Should not be able to instantiate interface directly
+    with pytest.raises(TypeError):
+        BaseWorkerGroup()
+
+    # Should not be able to create incomplete implementation
+    class IncompleteWorkerGroup(BaseWorkerGroup):
+        def execute(self, func, *args, **kwargs):
+            pass
+
+        # Missing other abstract methods
+
+    with pytest.raises(TypeError):
+        IncompleteWorkerGroup()
+
+
+def test_real_implementations_inherit_interface():
+    """Smoke test that real WorkerGroup implementations inherit from interface."""
+    # Test inheritance
+    assert issubclass(V1WorkerGroup, BaseWorkerGroup)
+    assert issubclass(V2WorkerGroup, BaseWorkerGroup)
+
+    # Test that all abstract methods are implemented
+    # If any abstract methods are missing, __abstractmethods__ will be non-empty
+    assert (
+        len(V1WorkerGroup.__abstractmethods__) == 0
+    ), f"V1 WorkerGroup missing abstract methods: {V1WorkerGroup.__abstractmethods__}"
+    assert (
+        len(V2WorkerGroup.__abstractmethods__) == 0
+    ), f"V2 WorkerGroup missing abstract methods: {V2WorkerGroup.__abstractmethods__}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

python/ray/train/tests/test_checkpoint.py

@@ -0,0 +1,232 @@
+import os
+from pathlib import Path
+
+import pyarrow.fs
+import pytest
+
+import ray
+from ray.train._checkpoint import (
+    _CHECKPOINT_TEMP_DIR_PREFIX,
+    _METADATA_FILE_NAME,
+    Checkpoint,
+    _get_del_lock_path,
+    _list_existing_del_locks,
+)
+from ray.train._internal.storage import _exists_at_fs_path, _upload_to_fs_path
+from ray.train.tests.test_new_persistence import _create_mock_custom_fs
+
+_CHECKPOINT_CONTENT_FILE = "dummy.txt"
+
+
+@pytest.fixture(params=["local", "mock", "custom_fs"])
+def checkpoint(request, tmp_path):
+    """Fixture that sets up a checkpoint on different filesystems."""
+    checkpoint_fs_type = request.param
+
+    checkpoint_path = tmp_path / "ckpt_dir"
+    checkpoint_path.mkdir(exist_ok=True)
+    (checkpoint_path / _CHECKPOINT_CONTENT_FILE).write_text("dummy")
+
+    if checkpoint_fs_type == "local":
+        yield Checkpoint.from_directory(str(checkpoint_path))
+    elif checkpoint_fs_type == "mock":
+        _checkpoint = Checkpoint(path="mock:///mock_bucket/ckpt_dir")
+        _upload_to_fs_path(
+            local_path=str(checkpoint_path),
+            fs=_checkpoint.filesystem,
+            fs_path=_checkpoint.path,
+        )
+        # The "mock://" URI doesn't persist across different instances of
+        # the pyarrow.fs.MockFileSystem, so we must make sure to return
+        # the checkpoint with the same filesystem instance that we uploaded
+        # some mock content to.
+        yield _checkpoint
+    elif checkpoint_fs_type == "custom_fs":
+        custom_storage_fs = _create_mock_custom_fs(tmp_path / "custom_fs")
+        _upload_to_fs_path(
+            local_path=str(checkpoint_path),
+            fs=custom_storage_fs,
+            fs_path="mock_bucket/ckpt_dir",
+        )
+        yield Checkpoint(path="mock_bucket/ckpt_dir", filesystem=custom_storage_fs)
+
+
+def test_to_directory(checkpoint: Checkpoint):
+    checkpoint_path = Path(checkpoint.to_directory())
+
+    assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+    assert _CHECKPOINT_TEMP_DIR_PREFIX in checkpoint_path.name
+
+
+def test_to_directory_with_user_specified_path(checkpoint: Checkpoint, tmp_path):
+    # Test with a string
+    checkpoint_path = Path(checkpoint.to_directory(str(tmp_path / "special_dir")))
+    assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+    assert checkpoint_path.name == "special_dir"
+
+    # Test with a PathLike
+    checkpoint_path = Path(checkpoint.to_directory(tmp_path / "special_dir"))
+    assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+    assert checkpoint_path.name == "special_dir"
+
+
+def test_multiprocess_to_directory(checkpoint: Checkpoint):
+    """Test the case where multiple processes are trying to checkpoint.
+
+    Only one process should download the checkpoint, and the others should
+    wait until it's finished. In the end, a single checkpoint dir should be
+    shared by all processes.
+    """
+    if checkpoint.filesystem.type_name == "mock":
+        pytest.skip("Mock filesystem cannot be pickled for use with Ray.")
+
+    @ray.remote
+    def download_checkpoint(checkpoint: Checkpoint) -> str:
+        return checkpoint.to_directory()
+
+    paths = [ray.get(download_checkpoint.remote(checkpoint)) for _ in range(5)]
+    # Check that all the paths are the same (no duplicates).
+    assert len(set(paths)) == 1
+
+
+def test_as_directory(checkpoint: Checkpoint):
+    with checkpoint.as_directory() as checkpoint_path:
+        checkpoint_path = Path(checkpoint_path)
+        assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+
+        if isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
+            # We should have directly returned the local path
+            assert str(checkpoint_path) == checkpoint.path
+        else:
+            # We should have downloaded to a temp dir.
+            assert _CHECKPOINT_TEMP_DIR_PREFIX in checkpoint_path.name
+
+    if isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
+        # Checkpoint should not be deleted, if we directly gave the local path.
+        assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
+    else:
+        # Should have been deleted, if the checkpoint downloaded to a temp dir.
+        assert not checkpoint_path.exists()
+
+
+def test_multiprocess_as_directory(checkpoint: Checkpoint, monkeypatch):
+    """Tests that deletion lock files are created and deleted correctly,
+    when multiple processes are accessing a checkpoint with `as_directory`"""
+    # If this is pointing to a local path, no lockfiles will be created.
+    is_local_checkpoint = isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem)
+
+    monkeypatch.setattr(os, "getpid", lambda: 11111)
+
+    with checkpoint.as_directory() as checkpoint_dir_1:
+        lock_file_1 = Path(_get_del_lock_path(checkpoint_dir_1))
+
+        # Pretend that the 2nd `as_directory` is called from a different process.
+        monkeypatch.setattr(os, "getpid", lambda: 22222)
+
+        with checkpoint.as_directory() as checkpoint_dir_2:
+            lock_file_2 = Path(_get_del_lock_path(checkpoint_dir_2))
+
+            # Should point both processes to the same canonical temp directory.
+            assert checkpoint_dir_1 == checkpoint_dir_2
+
+            if is_local_checkpoint:
+                # Check that 2 different lock files were created.
+                assert not lock_file_1.exists()
+                assert not lock_file_2.exists()
+            else:
+                # Check that 2 different lock files were created.
+                assert lock_file_1.exists()
+                assert lock_file_2.exists()
+
+                assert len(_list_existing_del_locks(checkpoint_dir_1)) == 2
+
+        if not is_local_checkpoint:
+            # Check that the 2nd lock file was deleted.
+            assert len(_list_existing_del_locks(checkpoint_dir_1)) == 1
+            assert not lock_file_2.exists()
+
+    if not is_local_checkpoint:
+        # Check that the 1st lock file was deleted.
+        assert len(_list_existing_del_locks(checkpoint_dir_1)) == 0
+        assert not lock_file_1.exists()
+
+        # Check that the temp checkpoint directory was deleted.
+        assert not Path(checkpoint_dir_1).exists()
+
+
+def test_as_directory_lock_cleanup(checkpoint: Checkpoint):
+    """Errors when accessing a checkpoint with `as_directory`
+    shouldn't leave behind lock files.
+    """
+    with pytest.raises(RuntimeError):
+        with checkpoint.as_directory() as checkpoint_dir:
+            raise RuntimeError
+
+    assert not _list_existing_del_locks(checkpoint_dir)
+
+    is_local_checkpoint = isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem)
+    if not is_local_checkpoint:
+        assert not Path(checkpoint_dir).exists()
+
+
+def test_as_directory_download_error(checkpoint: Checkpoint, monkeypatch):
+    """Errors during a checkpoint download should be raised directly when accessing
+    it with the `as_directory` context manager."""
+    if isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
+        pytest.skip(
+            "Local filesystem checkpoints don't download to a temp dir, so "
+            "there's no error handling to test."
+        )
+
+    error_text = "original error"
+
+    def to_directory_error(*args, **kwargs):
+        raise RuntimeError(error_text)
+
+    monkeypatch.setattr(checkpoint, "to_directory", to_directory_error)
+
+    with pytest.raises(RuntimeError, match=error_text):
+        with checkpoint.as_directory() as _:
+            pass
+
+
+def test_metadata(checkpoint: Checkpoint):
+    assert checkpoint.get_metadata() == {}
+
+    # No metadata file by default.
+    assert not _exists_at_fs_path(
+        fs=checkpoint.filesystem,
+        fs_path=str(Path(checkpoint.path) / _METADATA_FILE_NAME),
+    )
+
+    checkpoint.update_metadata({"foo": "bar"})
+    assert checkpoint.get_metadata() == {"foo": "bar"}
+
+    checkpoint.update_metadata({"foo": "baz"})
+    assert checkpoint.get_metadata() == {"foo": "baz"}
+
+    checkpoint.update_metadata({"x": 1})
+    assert checkpoint.get_metadata() == {"foo": "baz", "x": 1}
+
+    # Set metadata completely resets the metadata.
+    checkpoint.set_metadata({"y": [1, 2, 3]})
+    assert checkpoint.get_metadata() == {"y": [1, 2, 3]}
+
+    # There should be a metadata file in the checkpoint directory.
+    assert _exists_at_fs_path(
+        fs=checkpoint.filesystem,
+        fs_path=str(Path(checkpoint.path) / _METADATA_FILE_NAME),
+    )
+
+    # Non JSON serializable metadata should raise an error.
+    class Test:
+        pass
+
+    with pytest.raises(TypeError):
+        checkpoint.set_metadata({"non_json_serializable": Test()})
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/train/tests/test_checkpoint_manager.py

@@ -0,0 +1,239 @@
+import random
+from pathlib import Path
+from typing import List
+
+import pytest
+
+from ray.train import Checkpoint, CheckpointConfig
+from ray.train._internal.checkpoint_manager import _CheckpointManager, _TrainingResult
+from ray.train.constants import TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE
+
+
+@pytest.fixture
+def checkpoint_paths(tmp_path):
+    checkpoint_paths = []
+    for i in range(10):
+        checkpoint_path = tmp_path / f"ckpt_{i}"
+        checkpoint_path.mkdir()
+        (checkpoint_path / "dummy.txt").write_text(f"{i}")
+        checkpoint_paths.append(checkpoint_path)
+
+    yield [str(path) for path in checkpoint_paths]
+
+
+def test_unlimited_checkpoints(checkpoint_paths: List[str]):
+    manager = _CheckpointManager(checkpoint_config=CheckpointConfig(num_to_keep=None))
+
+    for i in range(10):
+        manager.register_checkpoint(
+            _TrainingResult(
+                checkpoint=Checkpoint.from_directory(checkpoint_paths[i]),
+                metrics={"iter": i},
+            )
+        )
+
+    assert len(manager.best_checkpoint_results) == 10
+
+
+def test_limited_checkpoints(checkpoint_paths: List[str]):
+    manager = _CheckpointManager(checkpoint_config=CheckpointConfig(num_to_keep=2))
+
+    for i in range(10):
+        manager.register_checkpoint(
+            _TrainingResult(
+                checkpoint=Checkpoint.from_directory(checkpoint_paths[i]),
+                metrics={"iter": i},
+            )
+        )
+
+    assert len(manager.best_checkpoint_results) == 2
+
+    # Keep the latest checkpoints if no metric is given.
+    assert {
+        tracked_checkpoint.metrics["iter"]
+        for tracked_checkpoint in manager.best_checkpoint_results
+    } == {8, 9}
+
+    # The first 8 checkpoints should be deleted.
+    for i in range(8):
+        assert not Path(checkpoint_paths[i]).exists()
+
+    assert Path(checkpoint_paths[8]).exists()
+    assert Path(checkpoint_paths[9]).exists()
+
+
+@pytest.mark.parametrize("order", ["min", "max"])
+def test_keep_checkpoints_by_score(order, checkpoint_paths):
+    num_to_keep = 2
+    score_attribute = "score"
+
+    manager = _CheckpointManager(
+        checkpoint_config=CheckpointConfig(
+            num_to_keep=num_to_keep,
+            checkpoint_score_attribute=score_attribute,
+            checkpoint_score_order=order,
+        )
+    )
+
+    scores = []
+    for i in range(10):
+        score = random.random()
+        manager.register_checkpoint(
+            _TrainingResult(
+                checkpoint=Checkpoint.from_directory(checkpoint_paths[i]),
+                metrics={"iter": i, score_attribute: score},
+            )
+        )
+        scores.append(score)
+
+    sorted_scores = sorted(scores, reverse=order == "max")
+    assert set(sorted_scores[:num_to_keep]) == {
+        tracked_checkpoint.metrics[score_attribute]
+        for tracked_checkpoint in manager.best_checkpoint_results
+    }
+
+    # Make sure the bottom checkpoints are deleted.
+    best_checkpoint_iters = {
+        tracked_checkpoint.metrics["iter"]
+        for tracked_checkpoint in manager.best_checkpoint_results
+    }
+    for i, checkpoint_path in enumerate(checkpoint_paths):
+        if i in best_checkpoint_iters or i == 9:
+            # The checkpoint should only exist if it's one of the top K or the latest.
+            assert Path(checkpoint_path).exists()
+        else:
+            assert not Path(checkpoint_path).exists()
+
+
+def test_keep_latest_checkpoint(checkpoint_paths):
+    manager = _CheckpointManager(
+        checkpoint_config=CheckpointConfig(
+            num_to_keep=2,
+            checkpoint_score_attribute="score",
+            checkpoint_score_order="max",
+        )
+    )
+
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[0]),
+            metrics={"score": 3.0},
+        )
+    )
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[1]),
+            metrics={"score": 2.0},
+        )
+    )
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[2]),
+            metrics={"score": 1.0},
+        )
+    )
+
+    assert len(manager.best_checkpoint_results) == 2
+
+    # The latest checkpoint with the lowest score should not be deleted yet.
+    assert manager.latest_checkpoint_result.metrics["score"] == 1.0
+
+    # The latest checkpoint with the lowest score should not be deleted yet.
+    assert Path(checkpoint_paths[2]).exists()
+
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[3]),
+            metrics={"score": 0.0},
+        )
+    )
+    # A newer checkpoint came in. Even though the new one has a lower score, there are
+    # already num_to_keep better checkpoints, so the previous one should be deleted.
+    assert not Path(checkpoint_paths[2]).exists()
+
+    # Quick sanity check to make sure that the new checkpoint is kept.
+    assert manager.latest_checkpoint_result.metrics["score"] == 0.0
+    assert Path(checkpoint_paths[3]).exists()
+
+    # The original 2 checkpoints should still exist
+    assert Path(checkpoint_paths[0]).exists()
+    assert Path(checkpoint_paths[1]).exists()
+
+
+@pytest.mark.parametrize(
+    "metrics",
+    [
+        {"nested": {"sub": {"attr": 5}}},
+        {"nested": {"sub/attr": 5}},
+        {"nested/sub": {"attr": 5}},
+        {"nested/sub/attr": 5},
+    ],
+)
+def test_nested_get_checkpoint_score(metrics):
+    manager = _CheckpointManager(
+        checkpoint_config=CheckpointConfig(
+            num_to_keep=2,
+            checkpoint_score_attribute="nested/sub/attr",
+            checkpoint_score_order="max",
+        )
+    )
+
+    tracked_checkpoint = _TrainingResult(checkpoint=None, metrics=metrics)
+    assert manager._get_checkpoint_score(tracked_checkpoint) == (True, 5.0)
+
+
+@pytest.mark.parametrize("has_score_attr", [True, False])
+def test_only_store_score_attr(has_score_attr, checkpoint_paths, monkeypatch):
+    monkeypatch.setenv(TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE, "1")
+
+    # Set up CheckpointManager.
+    if has_score_attr:
+        checkpoint_config = CheckpointConfig(
+            num_to_keep=None,
+            checkpoint_score_attribute="score",
+            checkpoint_score_order="max",
+        )
+    else:
+        checkpoint_config = CheckpointConfig(num_to_keep=None)
+    manager = _CheckpointManager(checkpoint_config=checkpoint_config)
+
+    # Ensure we insert TrainingResults with score in the right order.
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[0]),
+            metrics={"score": 3.0},
+        )
+    )
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[1]),
+            metrics={"score": 1.0, "another_unsaved_metric": 6.0},
+        )
+    )
+    manager.register_checkpoint(
+        _TrainingResult(
+            checkpoint=Checkpoint.from_directory(checkpoint_paths[2]),
+            metrics={"another_unsaved_metric": 1.0},
+        )
+    )
+    assert len(manager.best_checkpoint_results) == 3
+    if has_score_attr:
+        assert manager.best_checkpoint_results[0].metrics == {"score": 1.0}
+        assert manager.best_checkpoint_results[0].checkpoint.path == checkpoint_paths[1]
+        assert manager.best_checkpoint_results[1].metrics == {"score": 3.0}
+        assert manager.best_checkpoint_results[1].checkpoint.path == checkpoint_paths[0]
+        assert manager.best_checkpoint_results[2].metrics == {}
+        assert manager.best_checkpoint_results[2].checkpoint.path == checkpoint_paths[2]
+    else:
+        assert manager.best_checkpoint_results[0].metrics == {}
+        assert manager.best_checkpoint_results[0].checkpoint.path == checkpoint_paths[0]
+        assert manager.best_checkpoint_results[1].metrics == {}
+        assert manager.best_checkpoint_results[1].checkpoint.path == checkpoint_paths[1]
+        assert manager.best_checkpoint_results[2].metrics == {}
+        assert manager.best_checkpoint_results[2].checkpoint.path == checkpoint_paths[2]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/train/tests/test_data_parallel_trainer.py

@@ -0,0 +1,372 @@
+import os
+import time
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray import train, tune
+from ray._common.utils import RESOURCE_CONSTRAINT_PREFIX
+from ray.train import RunConfig, ScalingConfig
+from ray.train._internal.backend_executor import BackendExecutor
+from ray.train._internal.worker_group import WorkerGroup
+from ray.train.backend import Backend, BackendConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.utils import _in_ray_train_worker
+from ray.tune.callback import Callback
+from ray.tune.tune_config import TuneConfig
+from ray.tune.tuner import Tuner
+from ray.util.accelerators import NVIDIA_A100
+
+
+@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()
+
+
+@pytest.fixture
+def ray_start_4_cpus_4_gpus_4_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=4, resources={"extra": 4})
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_4_cpus_4_gpus_4_a100():
+    address_info = ray.init(
+        num_cpus=4,
+        num_gpus=4,
+        resources={f"{RESOURCE_CONSTRAINT_PREFIX}{NVIDIA_A100}": 4},
+    )
+    yield address_info
+    ray.shutdown()
+
+
+def gen_execute_single_async_special(special_f):
+    def execute_single_async_special(self, i, f, *args, **kwargs):
+        assert len(self.workers) == 2
+        if i == 0 and hasattr(self, "should_fail") and self.should_fail:
+            kwargs["train_func"] = special_f
+        return (
+            self.workers[i]
+            .actor._RayTrainWorker__execute.options(name=f.__name__)
+            .remote(f, *args, **kwargs)
+        )
+
+    return execute_single_async_special
+
+
+def gen_new_backend_executor(special_f):
+    """Returns a BackendExecutor that runs special_f on worker 0 once."""
+
+    class TestBackendExecutor(BackendExecutor):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self._has_failed = False
+
+        def start_training(self, *args, **kwargs):
+            special_execute = gen_execute_single_async_special(special_f)
+            if not self._has_failed:
+                self.worker_group.should_fail = True
+                self._has_failed = True
+            else:
+                self.worker_group.should_fail = False
+            with patch.object(WorkerGroup, "execute_single_async", special_execute):
+                super().start_training(*args, **kwargs)
+
+    return TestBackendExecutor
+
+
+class CaptureReportCallback(Callback):
+    def __init__(self):
+        self.result_list = []
+
+    def on_trial_result(self, iteration, trials, trial, result, **info):
+        self.result_list.append(result)
+
+
+scale_config = ScalingConfig(num_workers=2)
+
+
+def test_fit_train(ray_start_4_cpus):
+    def train_func():
+        train.report({"loss": 1})
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func, scaling_config=scale_config
+    )
+    assert trainer.fit().metrics["loss"] == 1
+
+
+def test_scaling_config(ray_start_4_cpus):
+    def train_func():
+        assert ray.available_resources()["CPU"] == 1
+        train.report({"loss": 1})
+
+    assert ray.available_resources()["CPU"] == 4
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=2)
+    )
+    trainer.fit()
+
+
+def test_fit_train_config(ray_start_4_cpus):
+    def train_func(config):
+        train.report({"loss": config["x"]})
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func,
+        scaling_config=scale_config,
+        train_loop_config={"x": 100},
+    )
+    assert trainer.fit().metrics["loss"] == 100
+
+
+def test_datasets(ray_start_4_cpus):
+    num_train_data = 10
+    num_val_data = 6
+
+    train_dataset = ray.data.range(num_train_data)
+    val_dataset = ray.data.range(num_val_data)
+
+    def get_dataset():
+        # Train dataset should be sharded.
+        train_dataset = train.get_dataset_shard("train")
+        train_ds_count = len(list(train_dataset.iter_rows()))
+        assert train_ds_count == num_train_data / scale_config.num_workers
+        # All other datasets should not be sharded.
+        val_dataset = train.get_dataset_shard("val")
+        val_ds_count = len(list(val_dataset.iter_rows()))
+        assert val_ds_count == num_val_data / scale_config.num_workers
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=get_dataset,
+        scaling_config=scale_config,
+        datasets={"train": train_dataset, "val": val_dataset},
+    )
+    trainer.fit()
+
+
+def test_invalid_train_loop():
+    def train_loop(config, extra_arg):
+        pass
+
+    with pytest.raises(ValueError):
+        DataParallelTrainer(train_loop_per_worker=train_loop)
+
+
+def test_bad_return_in_train_loop(ray_start_4_cpus):
+    """Test to check if returns from train loop are discarded."""
+
+    # Simulates what happens with eg. torch models
+    class FailOnUnpickle:
+        def __reduce__(self):
+            raise RuntimeError("Failing")
+
+    def train_loop(config):
+        train.report({"loss": 1})
+        return FailOnUnpickle()
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_loop, scaling_config=scale_config
+    )
+
+    # No exception should happen here
+    trainer.fit()
+
+
+def test_tune(ray_start_4_cpus):
+    def train_func(config):
+        train.report({"loss": config["x"]})
+
+    trainer = DataParallelTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config={"x": 100},
+        scaling_config=scale_config,
+    )
+
+    tuner = Tuner(
+        trainer,
+        param_space={"train_loop_config": {"x": tune.choice([200, 300])}},
+        tune_config=TuneConfig(num_samples=2),
+    )
+    result_grid = tuner.fit()
+    assert result_grid[0].metrics["loss"] in [200, 300]
+
+    # Make sure original Trainer is not affected.
+    assert trainer._train_loop_config["x"] == 100
+
+
+def test_fast_slow(ray_start_4_cpus):
+    def train_func():
+        for i in range(2):
+            with create_dict_checkpoint({"epoch": i}) as checkpoint:
+                train.report(dict(index=i), checkpoint=checkpoint)
+
+    def train_slow():
+        for i in range(2):
+            with create_dict_checkpoint({"epoch": i}) as checkpoint:
+                train.report(dict(index=i), checkpoint=checkpoint)
+            time.sleep(5)
+
+    new_backend_executor_cls = gen_new_backend_executor(train_slow)
+    callback = CaptureReportCallback()
+
+    class DataParallelTrainerPatched(DataParallelTrainer):
+        _backend_executor_cls = new_backend_executor_cls
+
+    trainer = DataParallelTrainerPatched(
+        train_func,
+        scaling_config=scale_config,
+        run_config=RunConfig(callbacks=[callback]),
+    )
+    results = trainer.fit()
+
+    assert load_dict_checkpoint(results.checkpoint)["epoch"] == 1
+
+    result_list = callback.result_list
+    assert len(result_list) == 2
+
+
+def test_mismatch_report(ray_start_4_cpus):
+    def train_func():
+        for _ in range(2):
+            train.report(dict(loss=1))
+
+    def train_mismatch():
+        train.report(dict(loss=1))
+
+    new_backend_executor_cls = gen_new_backend_executor(train_mismatch)
+
+    class DataParallelTrainerPatched(DataParallelTrainer):
+        _backend_executor_cls = new_backend_executor_cls
+
+    trainer = DataParallelTrainerPatched(
+        train_func,
+        scaling_config=scale_config,
+    )
+    with pytest.raises(RuntimeError):
+        trainer.fit()
+
+
+def test_world_rank(ray_start_4_cpus, tmp_path):
+    def train_func():
+        world_rank = train.get_context().get_world_rank()
+        (tmp_path / f"{world_rank}").touch()
+        train.report(dict(world_rank=world_rank))
+
+    trainer = DataParallelTrainer(train_func, scaling_config=scale_config)
+    trainer.fit()
+
+    created_files = list(tmp_path.glob("*"))
+    assert len(created_files) == 2
+    assert {int(file.name) for file in created_files} == {0, 1}
+
+
+def test_gpu_requests(ray_start_4_cpus_4_gpus_4_extra, tmp_path):
+    def get_visible_devices_for_workers():
+        return [file.read_text() for file in tmp_path.glob("*")]
+
+    class CudaTestBackend(Backend):
+        share_cuda_visible_devices = True
+
+    class CudaTestConfig(BackendConfig):
+        @property
+        def backend_cls(self):
+            return CudaTestBackend
+
+    def get_resources():
+        cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+        world_rank = train.get_context().get_world_rank()
+        (tmp_path / f"{world_rank}").write_text(cuda_visible_devices)
+        train.report(dict(devices=cuda_visible_devices))
+
+    # 0 GPUs will be requested and should not raise an error.
+    trainer = DataParallelTrainer(
+        get_resources,
+        backend_config=CudaTestConfig(),
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
+    )
+    trainer.fit()
+    assert get_visible_devices_for_workers() == ["", ""]
+
+    # 1 GPU will be requested and should not raise an error.
+    trainer = DataParallelTrainer(
+        get_resources,
+        backend_config=CudaTestConfig(),
+        scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
+    )
+    trainer.fit()
+    visible_devices = get_visible_devices_for_workers()
+    # Sort the cuda visible devices to have exact match with expected result.
+    visible_devices = [",".join(sorted(r.split(","))) for r in visible_devices]
+    assert visible_devices == ["0,1", "0,1"]
+
+    # Partial GPUs should not raise an error.
+    trainer = DataParallelTrainer(
+        get_resources,
+        backend_config=CudaTestConfig(),
+        scaling_config=ScalingConfig(
+            num_workers=2, use_gpu=True, resources_per_worker={"GPU": 0.1}
+        ),
+    )
+    trainer.fit()
+    visible_devices = get_visible_devices_for_workers()
+    assert visible_devices == ["0", "0"]
+
+    # Multiple GPUs should not raise an error.
+    trainer = DataParallelTrainer(
+        get_resources,
+        backend_config=CudaTestConfig(),
+        scaling_config=ScalingConfig(
+            num_workers=2, use_gpu=True, resources_per_worker={"GPU": 2}
+        ),
+    )
+    trainer.fit()
+    visible_devices = get_visible_devices_for_workers()
+    # Sort the cuda visible devices to have exact match with expected result.
+    visible_devices = [",".join(sorted(r.split(","))) for r in visible_devices]
+    assert visible_devices == ["0,1,2,3", "0,1,2,3"]
+
+
+@pytest.mark.parametrize("accelerator_type", [NVIDIA_A100, None])
+def test_config_accelerator_type(ray_start_4_cpus_4_gpus_4_a100, accelerator_type):
+    def train_func():
+        # Ensure all workers are scheduled on nodes with specified accelerators
+        assigned_resources = ray.get_runtime_context().get_assigned_resources()
+        if accelerator_type:
+            accelerator_key = f"{RESOURCE_CONSTRAINT_PREFIX}{accelerator_type}"
+            assert accelerator_key in assigned_resources
+
+    trainer = DataParallelTrainer(
+        train_func,
+        scaling_config=ScalingConfig(
+            num_workers=4,
+            use_gpu=True,
+            accelerator_type=accelerator_type,
+        ),
+    )
+    trainer.fit()
+
+
+def test_in_ray_train_worker(ray_start_4_cpus):
+    assert not _in_ray_train_worker()
+
+    def train_fn():
+        assert _in_ray_train_worker()
+
+    trainer = DataParallelTrainer(train_fn)
+    trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))