chore: import upstream snapshot with attribution

python/ray/train/v2/_internal/callbacks/__init__.py

@@ -0,0 +1,16 @@
+from .accelerators import AcceleratorSetupCallback
+from .backend_setup import BackendSetupCallback
+from .datasets import DatasetsCallback
+from .state_manager import StateManagerCallback
+from .working_dir_setup import WorkingDirectorySetupCallback
+
+__all__ = [
+    "AcceleratorSetupCallback",
+    "BackendSetupCallback",
+    "DatasetsCallback",
+    "StateManagerCallback",
+    "WorkingDirectorySetupCallback",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/v2/_internal/callbacks/accelerators.py

@@ -0,0 +1,160 @@
+import logging
+import os
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Dict, List
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._private.accelerators.nvidia_gpu import CUDA_VISIBLE_DEVICES_ENV_VAR
+from ray._private.ray_constants import env_bool
+from ray.train import BackendConfig
+from ray.train.constants import ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV
+from ray.train.v2._internal.execution.callback import WorkerGroupCallback
+from ray.train.v2._internal.execution.worker_group import ActorMetadata
+from ray.train.v2.api.config import ScalingConfig
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.execution.worker_group.worker import Worker
+
+logger = logging.getLogger(__name__)
+
+
+class AcceleratorSetupCallback(WorkerGroupCallback):
+    """Perform accelerator setup for workers.
+
+    For example, this callback can be used to share CUDA_VISIBLE_DEVICES
+    among workers on the same node.
+    """
+
+    def __init__(self, backend_config: BackendConfig, scaling_config: ScalingConfig):
+        self._backend = backend_config.backend_cls()
+        self._scaling_config = scaling_config
+
+    def before_init_train_context(
+        self, workers: List["Worker"]
+    ) -> Dict[str, List[Any]]:
+        self._maybe_share_cuda_visible_devices(workers)
+        # TODO: Add support for sharing other accelerator resources.
+
+        return {}
+
+    def _maybe_share_cuda_visible_devices(self, workers: List["Worker"]):
+        """Set CUDA visible devices environment variables on workers."""
+        share_cuda_visible_devices_enabled = env_bool(
+            ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
+            self._backend.share_cuda_visible_devices,
+        )
+
+        if (
+            self._scaling_config._resources_per_worker_not_none.get("GPU", 0) > 0
+            and share_cuda_visible_devices_enabled
+        ):
+            _share_cuda_visible_devices(workers)
+
+
+def _share_cuda_visible_devices(workers: List["Worker"]):
+    """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"
+
+    Args:
+        workers: List of worker objects.
+    """
+    _share_accelerator_ids(workers, ray_constants.GPU, CUDA_VISIBLE_DEVICES_ENV_VAR)
+
+
+def _share_accelerator_ids(
+    workers: List["Worker"], accelerator_name: 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"
+        - Worker3: "0,1"
+
+    Args:
+        workers: List of worker objects.
+        accelerator_name: The name of the accelerator.
+        env_var: The name of the environment variable to set.
+    """
+    worker_metadatas = [worker.metadata for worker in workers]
+    visible_accelerator_ids_per_worker = _get_visible_accelerator_ids_per_worker(
+        worker_metadatas=worker_metadatas, accelerator_name=accelerator_name
+    )
+
+    def set_accelerator_ids(accelerator_ids):
+        os.environ[env_var] = accelerator_ids
+
+    futures = []
+    for rank, visible_accelerator_ids in enumerate(visible_accelerator_ids_per_worker):
+        futures.append(
+            workers[rank].execute_async(
+                set_accelerator_ids, accelerator_ids=visible_accelerator_ids
+            )
+        )
+    ray.get(futures)
+
+
+def _get_visible_accelerator_ids_per_worker(
+    worker_metadatas: List[ActorMetadata], accelerator_name: str
+) -> List[str]:
+    """Returns a list of comma-separated accelerator IDs visible to each worker.
+
+    All workers on a node should have the same set of visible accelerators,
+    which is the union of accelerator ids of the workers.
+
+    Args:
+        worker_metadatas: The actor metadata for each worker.
+        accelerator_name: The name of the accelerator resource to inspect.
+
+    Returns:
+        A list of comma-separated accelerator ID strings. This list is the
+        same length as the number of workers.
+    """
+    for metadata in worker_metadatas:
+        if accelerator_name not in metadata.accelerator_ids:
+            raise ValueError(
+                f"Accelerator '{accelerator_name}' is not available on all workers. "
+                f"Got these available accelerators instead: {metadata.accelerator_ids}"
+            )
+
+    node_id_to_accelerator_ids = defaultdict(set)
+
+    for metadata in worker_metadatas:
+        node_id_to_accelerator_ids[metadata.node_id].update(
+            metadata.accelerator_ids[accelerator_name]
+        )
+
+    visible_accelerator_ids_per_worker = []
+    for worker_id in range(len(worker_metadatas)):
+        node_id = worker_metadatas[worker_id].node_id
+        accelerator_ids = sorted(node_id_to_accelerator_ids[node_id])
+        all_resource_ids = ",".join([str(id) for id in accelerator_ids])
+        visible_accelerator_ids_per_worker.append(all_resource_ids)
+
+    return visible_accelerator_ids_per_worker

python/ray/train/v2/_internal/callbacks/backend_setup.py

@@ -0,0 +1,33 @@
+import logging
+
+from ray.exceptions import RayActorError, RayTaskError
+from ray.train.backend import BackendConfig
+from ray.train.v2._internal.execution.callback import (
+    ReplicaGroupCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.worker_group import (
+    ExecutionGroup,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class BackendSetupCallback(ReplicaGroupCallback, WorkerGroupCallback):
+    def __init__(self, backend_config: BackendConfig):
+        self._backend_config = backend_config
+        self._backend = backend_config.backend_cls()
+
+    def after_execution_group_start(self, execution_group: ExecutionGroup):
+        self._backend.on_start(execution_group, self._backend_config)
+        self._backend.on_training_start(execution_group, self._backend_config)
+
+    def before_execution_group_shutdown(self, execution_group: ExecutionGroup):
+        try:
+            self._backend.on_shutdown(execution_group, self._backend_config)
+        except (RayActorError, RayTaskError):
+            logger.warning(
+                "Graceful shutdown of backend failed. This is "
+                "expected if one of the workers has crashed.",
+                exc_info=True,
+            )

python/ray/train/v2/_internal/callbacks/datasets.py

@@ -0,0 +1,175 @@
+import copy
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import ray
+import ray.train
+from ray.train.v2._internal.data_integration.interfaces import (
+    DatasetShardMetadata,
+    DatasetShardProvider,
+    GenDataset,
+)
+from ray.train.v2._internal.execution.callback import WorkerGroupCallback
+from ray.train.v2._internal.execution.context import TrainRunContext
+from ray.train.v2._internal.execution.worker_group.worker_group import (
+    Worker,
+    WorkerGroup,
+    WorkerGroupContext,
+)
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+if TYPE_CHECKING:
+    from ray.data import DataIterator, Dataset, NodeIdStr
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+class RayDatasetShardProvider:
+    def __init__(
+        self,
+        datasets: Dict[str, GenDataset],
+        data_config: ray.train.DataConfig,
+        data_context: "DataContext",
+        world_size: int,
+        worker_node_ids: List["NodeIdStr"],
+    ):
+        from ray.train.v2._internal.data_integration.dataset_manager import (
+            DatasetManager,
+        )
+
+        self._dataset_names = set(datasets)
+        self._dataset_manager = (
+            ray.remote(DatasetManager)
+            .options(
+                num_cpus=0,
+                scheduling_strategy=NodeAffinitySchedulingStrategy(
+                    ray.get_runtime_context().get_node_id(), soft=False
+                ),
+            )
+            .remote(
+                datasets=datasets,
+                data_config=data_config,
+                data_context=data_context,
+                world_size=world_size,
+                worker_node_ids=worker_node_ids,
+            )
+        )
+        self._cached_dataset_shards: Dict[str, "DataIterator"] = {}
+
+    def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
+        dataset_name = dataset_info.dataset_name
+        if dataset_name not in self._dataset_names:
+            raise KeyError(
+                f"Dataset shard for '{dataset_name}' not found. "
+                "Please ensure that the dataset is passed through the Trainer `datasets` "
+                "argument."
+            )
+
+        if dataset_name not in self._cached_dataset_shards:
+            self._cached_dataset_shards[dataset_name] = ray.get(
+                self._dataset_manager.get_dataset_shard.remote(dataset_info)
+            )
+
+        return self._cached_dataset_shards[dataset_name]
+
+    def shutdown_data_executors(self) -> None:
+        """
+        Attempts to eagerly shutdown the data executors for datasets, freeing resources allocated to data execution.
+        """
+        try:
+            self._dataset_manager.shutdown_data_executors.remote()
+        except Exception:
+            logger.debug("Failed to invoke remote cleanup of Dataset Manager.")
+
+
+class DatasetsCallback(WorkerGroupCallback):
+    """A callback for managing Ray Datasets for the worker group."""
+
+    def __init__(
+        self,
+        train_run_context: TrainRunContext,
+        datasets: Dict[str, "Dataset"],
+    ):
+        self._datasets = datasets
+        self._data_config = copy.deepcopy(train_run_context.dataset_config)
+        self._scaling_config = train_run_context.scaling_config
+        self._dataset_shard_provider: Optional[RayDatasetShardProvider] = None
+
+        # Capture the current DataContext to propagate it to
+        # the Train workers later.
+        # The propagation works in the following way:
+        # 1. This callback is created when user create the Trainer.
+        # 2. Then this callback will be passed to the Controller actor.
+        # 3. Lastly, when the worker group is initialized, the Controller
+        #    will call the `after_worker_group_start` callback to propagate
+        #    the DataContext to Train workers.
+        from ray.data.context import DataContext
+
+        self._data_context = copy.deepcopy(DataContext.get_current())
+
+    def get_train_total_resources(
+        self, scaling_config: ray.train.ScalingConfig
+    ) -> Dict[str, float]:
+        """Return the resources reserved for training, so that Data can exclude
+        these resources logically from its available pool."""
+        if scaling_config.elasticity_enabled:
+            # If Train is running with a variable number of workers,
+            # we can't provide a fixed number of resources to exclude.
+            # Instead, Train and Data should coordinate via the autoscaling
+            # coordinator to allocate resources dynamically.
+            return {}
+        return scaling_config.total_resources
+
+    # --------------------------
+    # WorkerGroupCallback
+    # --------------------------
+
+    def before_init_train_context(
+        self, workers: List[Worker]
+    ) -> Dict[str, List[DatasetShardProvider]]:
+        world_size = len(workers)
+        worker_node_ids = [worker.metadata.node_id for worker in workers]
+        datasets = {k: v() if callable(v) else v for k, v in self._datasets.items()}
+
+        # TODO: Move this to the constructor.
+        # Notify the DataConfig about the total resources reserved for training.
+        total_train_resources = self.get_train_total_resources(self._scaling_config)
+        self._data_config.set_train_total_resources(
+            total_train_resources.get("CPU", 0), total_train_resources.get("GPU", 0)
+        )
+
+        self._dataset_shard_provider = RayDatasetShardProvider(
+            datasets=datasets,
+            data_config=self._data_config,
+            data_context=self._data_context,
+            world_size=world_size,
+            worker_node_ids=worker_node_ids,
+        )
+        return {"dataset_shard_provider": [self._dataset_shard_provider] * world_size}
+
+    def after_worker_group_start(self, worker_group: WorkerGroup):
+        # Propagate DataContext
+        from ray.data.context import DataContext
+
+        def _propagate_data_context(ctx: "DataContext"):
+            DataContext._set_current(ctx)
+
+        worker_group.execute(
+            _propagate_data_context,
+            self._data_context,
+        )
+
+    def after_worker_group_shutdown(
+        self, worker_group_context: WorkerGroupContext
+    ) -> None:
+        shard_provider = self._dataset_shard_provider
+        if shard_provider:
+            shard_provider.shutdown_data_executors()
+
+    def after_worker_group_abort(
+        self, worker_group_context: WorkerGroupContext
+    ) -> None:
+        shard_provider = self._dataset_shard_provider
+        if shard_provider:
+            shard_provider.shutdown_data_executors()

python/ray/train/v2/_internal/callbacks/env_callback.py

@@ -0,0 +1,39 @@
+import importlib
+import os
+from typing import List
+
+from ray.train.v2._internal.constants import RAY_TRAIN_CALLBACKS_ENV_VAR
+from ray.train.v2._internal.execution.callback import RayTrainCallback
+
+
+def _initialize_env_callbacks() -> List[RayTrainCallback]:
+    """Initialize callbacks from environment variable.
+
+    Returns:
+        List of callbacks initialized from environment variable.
+    """
+    callbacks = []
+    callbacks_str = os.environ.get(RAY_TRAIN_CALLBACKS_ENV_VAR, "")
+    if not callbacks_str:
+        return callbacks
+
+    for callback_path in callbacks_str.split(","):
+        callback_path = callback_path.strip()
+        if not callback_path:
+            continue
+
+        try:
+            module_path, class_name = callback_path.rsplit(".", 1)
+            module = importlib.import_module(module_path)
+            callback_cls = getattr(module, class_name)
+            if not issubclass(callback_cls, RayTrainCallback):
+                raise TypeError(
+                    f"Callback class '{callback_path}' must be a subclass of "
+                    f"RayTrainCallback, got {type(callback_cls).__name__}"
+                )
+            callback = callback_cls()
+            callbacks.append(callback)
+        except (ImportError, AttributeError, ValueError, TypeError) as e:
+            raise ValueError(f"Failed to import callback from '{callback_path}'") from e
+
+    return callbacks

python/ray/train/v2/_internal/callbacks/metrics.py

@@ -0,0 +1,124 @@
+from contextlib import contextmanager
+from typing import Dict, Optional
+
+import ray
+from ray.train.v2._internal.execution.callback import (
+    ControllerCallback,
+    TrainContextCallback,
+    WorkerCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.context import TrainRunContext, get_train_context
+from ray.train.v2._internal.execution.controller.state import (
+    TrainControllerState,
+    TrainControllerStateType,
+)
+from ray.train.v2._internal.metrics.base import Metric
+from ray.train.v2._internal.metrics.controller import ControllerMetrics
+from ray.train.v2._internal.metrics.worker import WorkerMetrics
+from ray.train.v2._internal.util import time_monotonic
+
+
+class ControllerMetricsCallback(ControllerCallback, WorkerGroupCallback):
+    """Callback that records controller-specific metrics."""
+
+    def after_controller_start(self, train_run_context: TrainRunContext):
+        """Initialize metrics after controller starts."""
+        self._run_name = train_run_context.get_run_config().name
+        self._run_id = train_run_context.run_id
+        self._metrics: Dict[str, Metric] = ControllerMetrics.get_controller_metrics(
+            self._run_name, self._run_id
+        )
+        # Record initial state
+        self._metrics[ControllerMetrics.CONTROLLER_STATE].record(
+            TrainControllerStateType.INITIALIZING
+        )
+
+    async def before_controller_shutdown(self):
+        """Shutdown metrics before controller shuts down."""
+        for metric in self._metrics.values():
+            metric.reset()
+
+    def after_controller_state_update(
+        self,
+        previous_state: TrainControllerState,
+        current_state: TrainControllerState,
+    ):
+        """Record state transitions after controller state updates."""
+        self._metrics[ControllerMetrics.CONTROLLER_STATE].record(
+            current_state._state_type
+        )
+
+    @contextmanager
+    def on_worker_group_start(self):
+        """Measure time taken to start worker group."""
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[ControllerMetrics.WORKER_GROUP_START_TOTAL_TIME_S].record(
+            elapsed_time_s
+        )
+
+    @contextmanager
+    def on_worker_group_shutdown(self):
+        """Measure time taken to shutdown worker group."""
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[ControllerMetrics.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S].record(
+            elapsed_time_s
+        )
+
+
+class WorkerMetricsCallback(WorkerCallback, TrainContextCallback):
+    """Callback that records worker-specific metrics."""
+
+    def __init__(self, train_run_context: TrainRunContext):
+        self._run_name = train_run_context.get_run_config().name
+        self._run_id = train_run_context.run_id
+        self._metrics: Optional[Dict[str, Metric]] = None
+
+    def after_init_train_context(self):
+        """Initialize metrics after train context is initialized."""
+        train_context = get_train_context()
+        core_context = ray.runtime_context.get_runtime_context()
+        world_rank = train_context.get_world_rank()
+        worker_actor_id = core_context.get_actor_id()
+        self._metrics = WorkerMetrics.get_worker_metrics(
+            self._run_name, self._run_id, world_rank, worker_actor_id
+        )
+
+    def before_worker_shutdown(self):
+        """Shutdown metrics before shutdown."""
+        if self._metrics:
+            for metric in self._metrics.values():
+                metric.reset()
+
+    @contextmanager
+    def on_report(self):
+        """
+        Context manager to measure the time taken to report a checkpoint to the storage.
+        """
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[WorkerMetrics.REPORT_TOTAL_BLOCKED_TIME_S].record(elapsed_time_s)
+
+    @contextmanager
+    def on_checkpoint_sync(self):
+        """Measure time spent in the cross-rank barrier that synchronizes the
+        checkpoint directory name across all workers."""
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[WorkerMetrics.CHECKPOINT_SYNC_TOTAL_TIME_S].record(elapsed_time_s)
+
+    @contextmanager
+    def on_checkpoint_transfer(self):
+        """Measure time spent transferring checkpoint files to storage."""
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[WorkerMetrics.CHECKPOINT_TRANSFER_TOTAL_TIME_S].record(
+            elapsed_time_s
+        )

python/ray/train/v2/_internal/callbacks/placement_group_callback.py

@@ -0,0 +1,131 @@
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import ray
+from ray.exceptions import RayActorError
+from ray.train.v2._internal.constants import GET_ACTOR_TIMEOUT_S
+from ray.train.v2._internal.execution.callback import (
+    ControllerCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.controller.placement_group_cleaner import (
+    PlacementGroupCleaner,
+)
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.execution.context import TrainRunContext
+    from ray.train.v2._internal.execution.worker_group import WorkerGroup
+
+logger = logging.getLogger(__name__)
+
+
+class PlacementGroupCleanerCallback(ControllerCallback, WorkerGroupCallback):
+    """Callback that manages a PlacementGroupCleaner for the training controller.
+
+    This callback ensures that placement groups are cleaned up even if the controller
+    dies ungracefully.
+    """
+
+    def __init__(
+        self,
+        check_interval_s: float = 1.0,
+        get_actor_timeout_s: float = GET_ACTOR_TIMEOUT_S,
+        stop_timeout: Optional[float] = None,
+    ):
+        """Initialize the callback.
+
+        Args:
+            check_interval_s: How often (in seconds) the cleaner should check
+                if the controller is still alive.
+            get_actor_timeout_s: How long to wait when calling the get actor state api.
+            stop_timeout: How long to wait for the cleaner to stop.
+        """
+        self._check_interval_s = check_interval_s
+        self._get_actor_timeout_s = get_actor_timeout_s
+        self._stop_timeout = stop_timeout
+        if self._stop_timeout is None:
+            self._stop_timeout = max(
+                2.0, self._check_interval_s * 2 + self._get_actor_timeout_s
+            )
+        self._cleaner: Optional[PlacementGroupCleaner] = None
+        self._controller_actor_id: Optional[str] = None
+
+    def after_controller_start(self, train_run_context: "TrainRunContext"):
+        """Launch the detached PlacementGroupCleaner actor and start monitoring."""
+
+        core_context = ray.runtime_context.get_runtime_context()
+        self._controller_actor_id = core_context.get_actor_id()
+        try:
+            # Launch the cleaner as a detached actor so it survives controller death
+            cleaner_actor_cls = ray.remote(num_cpus=0)(PlacementGroupCleaner)
+            self._cleaner = cleaner_actor_cls.options(
+                lifetime="detached",
+                get_if_exists=False,
+            ).remote(
+                controller_actor_id=self._controller_actor_id,
+                check_interval_s=self._check_interval_s,
+                get_actor_timeout_s=self._get_actor_timeout_s,
+                stop_timeout=self._stop_timeout,
+            )
+
+            logger.debug(
+                f"PlacementGroupCleaner launched for run_id={train_run_context.run_id}"
+            )
+        except Exception as e:
+            logger.warning(
+                f"Failed to launch PlacementGroupCleaner: {e}. "
+                "Placement groups may not be cleaned up if controller exits ungracefully."
+            )
+            self._cleaner = None
+            return
+
+        self._cleaner.start_monitoring.remote()
+
+    def after_worker_group_start(self, worker_group: "WorkerGroup"):
+        """Register the worker group's placement group with the cleaner.
+
+        This is called after a worker group is successfully started.
+        """
+        if not self._cleaner or not self._controller_actor_id:
+            logger.warning(
+                "PlacementGroupCleaner not available. "
+                "Placement groups may not be cleaned up if controller exits ungracefully."
+            )
+            return
+        worker_group_state = worker_group.get_worker_group_state()
+        placement_group = worker_group_state.placement_group_handle.placement_group
+
+        try:
+            ray.get(self._cleaner.register_placement_group.remote(placement_group))
+        except Exception as e:
+            logger.warning(
+                f"Failed to register placement group with cleaner: {e}. "
+                "Placement group may not be cleaned up if controller dies ungracefully."
+            )
+            return
+
+        logger.debug(
+            f"Registered placement group {placement_group.id} with PlacementGroupCleaner."
+        )
+
+    async def before_controller_shutdown(self):
+        self._stop_cleaner()
+
+    def before_controller_abort(self):
+        self._stop_cleaner()
+
+    def _stop_cleaner(self):
+        if not self._cleaner:
+            return
+
+        try:
+            # Stop the cleaner gracefully (it won't clean up the PG)
+            ray.get(self._cleaner.stop.remote(), timeout=self._stop_timeout)
+        except RayActorError:
+            logger.debug(
+                "PlacementGroupCleaner exited before stop completed; ignoring."
+            )
+        except Exception:
+            logger.exception("Failed to stop PlacementGroupCleaner gracefully.")
+        finally:
+            self._cleaner = None

python/ray/train/v2/_internal/callbacks/preemption_callback.py

@@ -0,0 +1,93 @@
+import logging
+import os
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import ray
+from ray.actor import ActorHandle
+from ray.train.v2._internal.constants import (
+    DEFAULT_PREEMPTION_POLL_INTERVAL_S,
+    PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
+)
+from ray.train.v2._internal.execution.callback import WorkerGroupCallback
+from ray.train.v2._internal.execution.preemption import PreemptionWatcher
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.execution.worker_group import (
+        WorkerGroup,
+        WorkerGroupContext,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+class PreemptionCallback(WorkerGroupCallback):
+    """Manages a :class:`PreemptionWatcher` across worker-group lifecycles.
+
+    Spawns a fresh watcher in :meth:`after_worker_group_start` and stops it on
+    every teardown path (shutdown and abort). Each worker group gets its own
+    watcher and failure-domain map, so elastic resizes and restarts never
+    leak stale state.
+    """
+
+    def __init__(self) -> None:
+        self._poll_interval_s: float = float(
+            os.getenv(
+                PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
+                str(DEFAULT_PREEMPTION_POLL_INTERVAL_S),
+            )
+        )
+        self._watcher: Optional[ActorHandle] = None
+
+    def after_worker_group_start(self, worker_group: "WorkerGroup") -> None:
+        # Tear down any watcher from a previous worker group first. Worker-group
+        # startup can fail after this hook without running the shutdown hook, so
+        # this also prevents leaking an orphaned watcher across a reschedule.
+        self._stop_watcher()
+
+        # These handles are captured once per worker-group start. With the
+        # standard backend, any worker replacement goes through a full worker
+        # group restart (this hook runs again with fresh handles), so they
+        # never go stale.
+        # TODO(lehui): refresh worker handles on in-place replica replacement
+        # when adding preemption support for replica groups (TorchFT).
+        node_to_ranks: Dict[str, List[int]] = {}
+        worker_actors_by_rank: Dict[int, ActorHandle] = {}
+        for w in worker_group.get_workers():
+            rank = w.distributed_context.world_rank
+            node_to_ranks.setdefault(w.metadata.node_id, []).append(rank)
+            worker_actors_by_rank[rank] = w.actor
+
+        watcher_cls = ray.remote(num_cpus=0, max_restarts=-1)(PreemptionWatcher)
+        self._watcher = watcher_cls.remote(
+            node_to_ranks=node_to_ranks,
+            poll_interval_s=self._poll_interval_s,
+            worker_actors_by_rank=worker_actors_by_rank,
+        )
+
+        logger.debug(
+            "PreemptionCallback: started watcher for %d node(s).",
+            len(node_to_ranks),
+        )
+
+    def before_worker_group_shutdown(self, worker_group: "WorkerGroup") -> None:
+        self._stop_watcher()
+
+    def after_worker_group_abort(
+        self, worker_group_context: "WorkerGroupContext"
+    ) -> None:
+        # abort() doesn't run the shutdown hook, so tear the watcher down here
+        # too — otherwise it keeps polling GCS until the cluster reaps it.
+        self._stop_watcher()
+
+    def _stop_watcher(self) -> None:
+        if self._watcher is None:
+            return
+        watcher = self._watcher
+        self._watcher = None
+        # Force-kill (non-blocking) rather than a synchronous graceful stop, so
+        # we never block the controller's event loop. The watcher's daemon poll
+        # thread dies with the actor process and holds no external resources.
+        try:
+            ray.kill(watcher)
+        except Exception:
+            logger.warning("Failed to kill PreemptionWatcher actor.", exc_info=True)

python/ray/train/v2/_internal/callbacks/state_manager.py

@@ -0,0 +1,221 @@
+import importlib
+import logging
+from typing import TYPE_CHECKING, Dict, Optional
+
+if TYPE_CHECKING:
+    from ray.data import Dataset
+
+import ray
+from ray.train.v2._internal.execution.callback import (
+    ControllerCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.context import TrainRunContext
+from ray.train.v2._internal.execution.controller.state import (
+    AbortedState,
+    ErroredState,
+    FinishedState,
+    ReschedulingState,
+    ResizingState,
+    RestartingState,
+    RunningState,
+    SchedulingState,
+    ShuttingDownState,
+    TrainControllerState,
+)
+from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
+    ResizeDecision,
+)
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroup,
+    WorkerGroupContext,
+    WorkerGroupState,
+)
+from ray.train.v2._internal.execution.worker_group.poll import WorkerGroupPollStatus
+from ray.train.v2._internal.logging.logging import (
+    get_train_application_controller_log_path,
+)
+from ray.train.v2._internal.state.state_manager import TrainStateManager
+from ray.train.v2._internal.util import TrainingFramework
+
+logger = logging.getLogger(__name__)
+
+
+def _get_framework_version(framework: Optional[TrainingFramework]):
+    versions = {}
+
+    try:
+        import ray
+
+        versions["ray"] = ray.__version__
+    except ImportError:
+        logger.warning("Failed to collect ray version on worker.")
+
+    if framework is None:
+        return versions
+
+    for module_name in framework.module_names():
+        try:
+            module = importlib.import_module(module_name)
+            versions[module_name] = module.__version__
+        except ModuleNotFoundError:
+            # Module is not installed, skip without recording a version.
+            continue
+        except Exception:
+            logger.warning(f"Failed to collect {module_name} version on worker.")
+            continue
+
+    return versions
+
+
+class StateManagerCallback(ControllerCallback, WorkerGroupCallback):
+    def __init__(self, datasets: Dict[str, "Dataset"]):
+        self._datasets = datasets
+
+    def after_controller_start(self, train_run_context: TrainRunContext):
+        self._state_manager = TrainStateManager()
+        self._run_name = train_run_context.get_run_config().name
+        self._run_id = train_run_context.run_id
+
+        # TODO: Should this be generated by the caller?
+        # NOTE: These must be called on the Controller.
+        #       The Callback is first initialized on the Driver.
+        core_context = ray.runtime_context.get_runtime_context()
+        self._job_id = core_context.get_job_id()
+        self._controller_actor_id = core_context.get_actor_id()
+        controller_log_file_path = get_train_application_controller_log_path()
+        self._state_manager.create_train_run(
+            id=self._run_id,
+            name=self._run_name,
+            job_id=self._job_id,
+            controller_actor_id=self._controller_actor_id,
+            controller_log_file_path=controller_log_file_path,
+            run_config=train_run_context.run_config,
+            train_loop_config=train_run_context.train_loop_config,
+            scaling_config=train_run_context.scaling_config,
+            backend_config=train_run_context.backend_config,
+            datasets=self._datasets,
+            dataset_config=train_run_context.dataset_config,
+        )
+
+    def after_controller_state_update(
+        self,
+        previous_state: TrainControllerState,
+        current_state: TrainControllerState,
+    ):
+        if previous_state._state_type == current_state._state_type:
+            return
+
+        logger.info(
+            f"[State Transition] {previous_state._state_type.state_name} -> "
+            f"{current_state._state_type.state_name}."
+        )
+
+        if isinstance(current_state, SchedulingState):
+            # TODO: This should probably always be ResizeDecision.
+            if isinstance(current_state.scaling_decision, ResizeDecision):
+                resize_decision = current_state.scaling_decision
+            else:
+                resize_decision = None
+
+            self._state_manager.update_train_run_scheduling(
+                run_id=self._run_id,
+                resize_decision=resize_decision,
+            )
+
+        elif isinstance(current_state, RunningState):
+            self._state_manager.update_train_run_running(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, RestartingState):
+            self._state_manager.update_train_run_restarting(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, ResizingState):
+            self._state_manager.update_train_run_resizing(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, ErroredState):
+            self._state_manager.update_train_run_errored(
+                run_id=self._run_id,
+                status_detail=str(current_state.training_failed_error),
+            )
+
+        elif isinstance(current_state, FinishedState):
+            self._state_manager.update_train_run_finished(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, AbortedState):
+            self._state_manager.update_train_run_aborted(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, ReschedulingState):
+            # substate of SchedulingState
+            pass
+
+        elif isinstance(current_state, ShuttingDownState):
+            # substate of RunningState
+            pass
+
+    def before_worker_group_start(self, worker_group_context: WorkerGroupContext):
+        self._state_manager.create_train_run_attempt(
+            run_id=self._run_id,
+            attempt_id=worker_group_context.run_attempt_id,
+            num_workers=worker_group_context.num_workers,
+            resources_per_worker=worker_group_context.resources_per_worker,
+        )
+
+    def after_worker_group_start(self, worker_group: WorkerGroup):
+        worker_group_context: WorkerGroupContext = (
+            worker_group.get_worker_group_context()
+        )
+        worker_group_state: WorkerGroupState = worker_group.get_worker_group_state()
+        self._state_manager.update_train_run_attempt_running(
+            run_id=self._run_id,
+            attempt_id=worker_group_context.run_attempt_id,
+            workers=worker_group_state.workers,
+        )
+
+        # Update train run framework version
+        framework = self._state_manager.get_train_run_framework(self._run_id)
+        framework_versions = worker_group.execute_single(
+            0, _get_framework_version, framework
+        )
+        self._state_manager.update_train_run_framework_versions(
+            run_id=self._run_id,
+            framework_versions=framework_versions,
+        )
+
+    def before_worker_group_shutdown(self, worker_group: WorkerGroup):
+        worker_group_context: WorkerGroupContext = (
+            worker_group.get_worker_group_context()
+        )
+        # TODO: Consider passing error reason directly to the callback.
+        # Something along the lines of:
+        #    WorkerGroup.shutdown(reason)
+        #    -> WorkerGroupCallback.before_worker_group_shutdown(reason)
+        worker_group_poll_status: Optional[
+            WorkerGroupPollStatus
+        ] = worker_group.get_latest_poll_status()
+        if worker_group_poll_status and worker_group_poll_status.errors:
+            self._state_manager.update_train_run_attempt_errored(
+                run_id=self._run_id,
+                attempt_id=worker_group_context.run_attempt_id,
+                status_detail=worker_group_poll_status.get_error_string(),
+            )
+        else:
+            self._state_manager.update_train_run_attempt_finished(
+                run_id=self._run_id,
+                attempt_id=worker_group_context.run_attempt_id,
+            )
+
+    def before_worker_group_abort(self, worker_group_context: WorkerGroupContext):
+        self._state_manager.update_train_run_attempt_aborted(
+            self._run_id,
+            worker_group_context.run_attempt_id,
+        )

python/ray/train/v2/_internal/callbacks/user_callback.py

@@ -0,0 +1,54 @@
+from typing import Any, Dict, List
+
+from ray.train.v2._internal.execution.callback import (
+    ReportCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.context import TrainRunContext
+from ray.train.v2._internal.execution.training_report import _TrainingReport
+from ray.train.v2._internal.execution.worker_group import WorkerGroupPollStatus
+from ray.train.v2.api.callback import UserCallback
+
+
+class UserCallbackHandler(WorkerGroupCallback, ReportCallback):
+    """Responsible for calling methods of subscribers implementing
+    the `UserCallback` interface.
+    """
+
+    def __init__(
+        self, user_callbacks: List[UserCallback], train_run_context: TrainRunContext
+    ):
+        self._user_callbacks = user_callbacks
+        self._train_run_context = train_run_context
+
+    # --------------------------
+    # ReportCallback
+    # --------------------------
+
+    def after_report(
+        self,
+        training_report: _TrainingReport,
+        metrics: List[Dict[str, Any]],
+    ):
+        for user_callback in self._user_callbacks:
+            user_callback.after_report(
+                run_context=self._train_run_context,
+                metrics=metrics,
+                checkpoint=training_report.checkpoint,
+            )
+
+    # --------------------------
+    # WorkerGroupCallback
+    # --------------------------
+
+    def after_worker_group_poll_status(
+        self, worker_group_status: WorkerGroupPollStatus
+    ):
+        if not worker_group_status.errors:
+            return
+
+        for user_callback in self._user_callbacks:
+            user_callback.after_exception(
+                run_context=self._train_run_context,
+                worker_exceptions=worker_group_status.errors,
+            )

python/ray/train/v2/_internal/callbacks/working_dir_setup.py

@@ -0,0 +1,31 @@
+import logging
+import os
+
+from ray.train.v2._internal.execution.callback import (
+    ReplicaGroupCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.context import get_train_context
+from ray.train.v2._internal.execution.worker_group import (
+    ExecutionGroup,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class WorkingDirectorySetupCallback(ReplicaGroupCallback, WorkerGroupCallback):
+    def after_execution_group_start(self, execution_group: ExecutionGroup):
+        """Shared logic for setting up the working directory on an execution group."""
+
+        def chdir_to_working_dir() -> None:
+            """Create the local working directory for the experiment."""
+            local_working_directory = (
+                get_train_context().get_storage().local_working_directory
+            )
+            os.makedirs(local_working_directory, exist_ok=True)
+            logger.debug(
+                f"Changing the working directory to: {local_working_directory}"
+            )
+            os.chdir(local_working_directory)
+
+        execution_group.execute(chdir_to_working_dir)

python/ray/train/v2/_internal/constants.py

@@ -0,0 +1,159 @@
+import os
+from typing import Dict, Optional
+
+from ray._common.constants import RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR
+from ray._private.ray_constants import env_bool, env_set_by_user
+
+# Unsupported configs can use this value to detect if the user has set it.
+_UNSUPPORTED = "UNSUPPORTED"
+_DEPRECATED = "DEPRECATED"
+
+# The name of the file that is used to validate the storage.
+VALIDATE_STORAGE_MARKER_FILENAME = ".validate_storage_marker"
+# The name of the file that is used to store the checkpoint manager snapshot.
+CHECKPOINT_MANAGER_SNAPSHOT_FILENAME = "checkpoint_manager_snapshot.json"
+
+AWS_RETRYABLE_TOKENS = (
+    "AWS Error SLOW_DOWN",
+    "AWS Error INTERNAL_FAILURE",
+    "AWS Error SERVICE_UNAVAILABLE",
+    "AWS Error NETWORK_CONNECTION",
+    "AWS Error UNKNOWN",
+)
+
+# -----------------------------------------------------------------------
+# Environment variables used in the controller, workers, and state actor.
+#
+# Be sure to update ENV_VARS_TO_PROPAGATE when adding new
+# environment variables in this section.
+# -----------------------------------------------------------------------
+
+# Polling interval for the Train controller.
+# This determines how many seconds the controller will wait between
+# polling the worker group for its status.
+HEALTH_CHECK_INTERVAL_S_ENV_VAR = "RAY_TRAIN_HEALTH_CHECK_INTERVAL_S"
+DEFAULT_HEALTH_CHECK_INTERVAL_S: float = 2.0
+
+# The time in seconds a worker health check must be hanging for
+# before the controller marks the worker as dead and handles the failure.
+WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR = "RAY_TRAIN_WORKER_HEALTH_CHECK_TIMEOUT_S"
+DEFAULT_WORKER_HEALTH_CHECK_TIMEOUT_S: float = 10 * 60
+
+# Timeout in seconds for the worker group to start.
+WORKER_GROUP_START_TIMEOUT_S_ENV_VAR = "RAY_TRAIN_WORKER_GROUP_START_TIMEOUT_S"
+DEFAULT_WORKER_GROUP_START_TIMEOUT_S: float = 60.0
+
+# Time in seconds for collective operations before raising a timeout error.
+COLLECTIVE_TIMEOUT_S_ENV_VAR = "RAY_TRAIN_COLLECTIVE_TIMEOUT_S"
+# NOTE: Default to no timeout to avoid introducing more timeouts for users to configure.
+# For example, users can already configure timeouts in torch distributed.
+DEFAULT_COLLECTIVE_TIMEOUT_S: Optional[float] = None
+# Interval in seconds to log a warning when waiting for a collective operation to complete.
+COLLECTIVE_WARN_INTERVAL_S_ENV_VAR = "RAY_TRAIN_COLLECTIVE_WARN_INTERVAL_S"
+DEFAULT_COLLECTIVE_WARN_INTERVAL_S: float = 60
+
+# Interval in seconds to log a warning when waiting for a checkpoint upload fn operation to complete.
+CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR = (
+    "RAY_TRAIN_CHECKPOINT_UPLOAD_WARN_INTERVAL_S"
+)
+DEFAULT_CHECKPOINT_UPLOAD_WARN_INTERVAL_S: float = 60
+
+# Feature flag for the preemption watcher. Default-on; provides a quick
+# rollback path if the watcher actor misbehaves in a cluster.
+ENABLE_PREEMPTION_WATCHER_ENV_VAR = "RAY_TRAIN_ENABLE_PREEMPTION_WATCHER"
+DEFAULT_ENABLE_PREEMPTION_WATCHER: bool = True
+
+# How often the preemption watcher polls Ray Core's drain state.
+PREEMPTION_POLL_INTERVAL_S_ENV_VAR = "RAY_TRAIN_PREEMPTION_POLL_INTERVAL_S"
+DEFAULT_PREEMPTION_POLL_INTERVAL_S: float = 5.0
+
+# Environment variable to enable the print function patching.
+ENABLE_PRINT_PATCH_ENV_VAR = "RAY_TRAIN_ENABLE_PRINT_PATCH"
+DEFAULT_ENABLE_PRINT_PATCH = True
+
+# V2 feature flag.
+V2_ENABLED_ENV_VAR = "RAY_TRAIN_V2_ENABLED"
+
+# Environment variables to enable/disable controller/worker structured logging.
+ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR = (
+    "RAY_TRAIN_ENABLE_CONTROLLER_STRUCTURED_LOGGING"
+)
+ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR = "RAY_TRAIN_ENABLE_WORKER_STRUCTURED_LOGGING"
+DEFAULT_ENABLE_CONTROLLER_LOGGING = True
+DEFAULT_ENABLE_WORKER_LOGGING = True
+
+# Environment variables to configure reconciliation interval for Train state actor.
+# This determines how many seconds the state actor will wait between
+# polling the controller for its status.
+ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR = (
+    "RAY_TRAIN_ENABLE_STATE_ACTOR_RECONCILIATION"
+)
+DEFAULT_ENABLE_STATE_ACTOR_RECONCILIATION = True
+STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR = (
+    "RAY_TRAIN_STATE_ACTOR_RECONCILIATION_INTERVAL_S"
+)
+DEFAULT_STATE_ACTOR_RECONCILIATION_INTERVAL_S: float = 30.0
+# TODO: `ray.util.state.api.get_actor` typically takes 10-50ms but can take longer
+# when there is high load on the cluster.
+GET_ACTOR_TIMEOUT_S: int = 10
+# GET_ACTOR_TIMEOUT_S * CONTROLLERS_TO_POLL_PER_ITERATION should be
+# way less than STATE_ACTOR_RECONCILIATION_INTERVAL_S to give the state actor
+# time to update live train run state.
+CONTROLLERS_TO_POLL_PER_ITERATION: int = 1
+
+# Environment variable for Train execution callbacks
+RAY_TRAIN_CALLBACKS_ENV_VAR = "RAY_TRAIN_CALLBACKS"
+
+# Ray Train does not warn by default when using blocking ray.get inside async actor.
+DEFAULT_RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_VALUE = "0"
+
+# torchft lighthouse address
+TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR = "TORCHFT_LIGHTHOUSE"
+
+# Environment variables to propagate from the driver to the controller,
+# and then from the controller to the workers.
+ENV_VARS_TO_PROPAGATE = {
+    V2_ENABLED_ENV_VAR,
+    HEALTH_CHECK_INTERVAL_S_ENV_VAR,
+    WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR,
+    WORKER_GROUP_START_TIMEOUT_S_ENV_VAR,
+    COLLECTIVE_TIMEOUT_S_ENV_VAR,
+    COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
+    CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR,
+    ENABLE_PRINT_PATCH_ENV_VAR,
+    ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR,
+    ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR,
+    ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
+    STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR,
+    RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR,
+    TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR,
+    ENABLE_PREEMPTION_WATCHER_ENV_VAR,
+    PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
+}
+
+
+# ------------------------------------------------------------
+# Environment variables used in the driver script only.
+# ------------------------------------------------------------
+
+# The environment variable to enable the Ray Train Metrics.
+METRICS_ENABLED_ENV_VAR = "RAY_TRAIN_METRICS_ENABLED"
+
+
+def is_v2_enabled() -> bool:
+    return env_bool(V2_ENABLED_ENV_VAR, True)
+
+
+def get_env_vars_to_propagate() -> Dict[str, str]:
+    """Returns a dictionary of environment variables that should be propagated
+    from the driver to the controller, and then from the controller
+    to each training worker.
+
+    This way, users only need to set environment variables in one place
+    when launching the script instead of needing to manually set a runtime environment.
+    """
+    env_vars = {}
+    for env_var in ENV_VARS_TO_PROPAGATE:
+        if env_set_by_user(env_var):
+            env_vars[env_var] = os.environ[env_var]
+    return env_vars

python/ray/train/v2/_internal/data_integration/dataset_manager.py

@@ -0,0 +1,174 @@
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Dict, List
+
+import ray
+from ray.exceptions import GetTimeoutError
+from ray.train.v2._internal.data_integration.interfaces import (
+    DatasetShardMetadata,
+    GenDataset,
+)
+
+if TYPE_CHECKING:
+    from ray.data import DataContext, DataIterator, Dataset, NodeIdStr
+
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetManager:
+    """Manages the dataset shards for datasets configured in the trainer."""
+
+    def __init__(
+        self,
+        datasets: Dict[str, GenDataset],
+        data_config: ray.train.DataConfig,
+        data_context: "DataContext",
+        world_size: int,
+        worker_node_ids: List["NodeIdStr"],
+    ):
+        self._datasets = datasets
+        self._data_config = data_config
+        self._datasets_to_split = (
+            set(self._datasets.keys())
+            if data_config._datasets_to_split == "all"
+            else set(data_config._datasets_to_split)
+        )
+        self._world_size = world_size
+        self._worker_node_ids = worker_node_ids
+        self._coordinator_actors: List[ray.actor.ActorHandle] = []
+
+        # Maps dataset name to a list of cached `DataIterator`s corresponding to
+        # Train worker ranks.
+        self._dataset_iterators: Dict[str, List["DataIterator"]] = {}
+
+        # A condition variable to synchronize the calls to the async `get_dataset_shard` method.
+        self._condition = asyncio.Condition()
+
+        from ray.data import DataContext
+
+        DataContext._set_current(data_context)
+
+    def _create_dataset_iterators(
+        self, dataset_info: DatasetShardMetadata, base_dataset: "Dataset"
+    ) -> List["DataIterator"]:
+        dataset_name = dataset_info.dataset_name
+
+        iterators_per_rank = self._data_config.configure(
+            datasets={dataset_name: base_dataset},
+            world_size=self._world_size,
+            worker_handles=None,
+            worker_node_ids=self._worker_node_ids,
+        )
+        assert len(iterators_per_rank) == self._world_size
+        # Convert the List[Dict[str, DataIterator]] to a List[DataIterator],
+        # since we only configured one dataset.
+        return [iterators_per_rank[i][dataset_name] for i in range(self._world_size)]
+
+    def _get_unsharded_dataset_iterator(
+        self, dataset_info: DatasetShardMetadata
+    ) -> "DataIterator":
+        """Returns the dataset iterator for a dataset that is excluded
+        from `DataConfig.datasets_to_split`.
+        Note that this method is NOT a barrier across workers and can be called
+        by any subset of workers and will return immediately.
+        """
+        dataset_name = dataset_info.dataset_name
+        world_rank = dataset_info.world_rank
+
+        if dataset_name not in self._dataset_iterators:
+            self._dataset_iterators[dataset_name] = self._create_dataset_iterators(
+                dataset_info, self._datasets[dataset_name]
+            )
+
+        return self._dataset_iterators[dataset_name][world_rank]
+
+    async def _get_sharded_dataset_iterator(
+        self, dataset_info: DatasetShardMetadata
+    ) -> "DataIterator":
+        """Returns the dataset iterator for a dataset that is included
+        in `DataConfig.datasets_to_split`.
+        Note that this method is a barrier across workers,
+        and all workers must call this method before training.
+        """
+        dataset_name = dataset_info.dataset_name
+        world_rank = dataset_info.world_rank
+
+        async with self._condition:
+            if dataset_name in self._dataset_iterators:
+                # If the dataset iterators have already been created, return the
+                # existing one.
+                iterator = self._dataset_iterators[dataset_name][world_rank]
+            elif world_rank == 0:
+                # In this case, the dataset iterators have not been created yet.
+                # The dataset only needs to be configured once globally for all workers.
+                # Do it only when the rank 0 worker calls this method.
+                iterators = self._create_dataset_iterators(
+                    dataset_info, self._datasets[dataset_name]
+                )
+                iterator = iterators[world_rank]
+
+                # Cache the split coordinators for resource cleanup.
+                from ray.data._internal.iterator.stream_split_iterator import (
+                    StreamSplitDataIterator,
+                )
+
+                if isinstance(iterator, StreamSplitDataIterator):
+                    self._coordinator_actors.append(iterator._coord_actor)
+
+                # Cache the dataset iterators for future use.
+                self._dataset_iterators[dataset_name] = iterators
+                self._condition.notify_all()
+            else:
+                # Wait for the dataset iterators to be created by the rank 0 worker.
+                await self._condition.wait_for(
+                    lambda: dataset_name in self._dataset_iterators
+                )
+                iterator = self._dataset_iterators[dataset_name][world_rank]
+        return iterator
+
+    async def get_dataset_shard(
+        self,
+        dataset_info: DatasetShardMetadata,
+    ) -> "DataIterator":
+        """Create and return the dataset shard iterator for a Ray Train worker's
+        call to `ray.train.get_dataset_shard`.
+
+        This method is a barrier that should be called by all Ray Train workers at once.
+        If the dataset iterators have already been created, return the existing ones.
+
+        Otherwise, create the dataset iterators and cache them.
+        Here's an example of how this method is used with 4 workers:
+        Rank 2 calls get_dataset_shard, waits on the condition variable.
+        Rank 1 calls get_dataset_shard, waits on the condition variable.
+        Rank 0 calls get_dataset_shard, creates the dataset iterators, caches them,
+        and notifies all workers hanging on the condition variable.
+        Rank 3 calls get_dataset_shard, returns the cached iterator.
+        """
+        dataset_name = dataset_info.dataset_name
+
+        if dataset_name in self._datasets_to_split:
+            return await self._get_sharded_dataset_iterator(dataset_info)
+        else:
+            return self._get_unsharded_dataset_iterator(dataset_info)
+
+    def shutdown_data_executors(self) -> None:
+        """
+        Attempts to shut down the data executors of each sharded dataset,
+        freeing resources allocated to data execution.
+
+        Note: The data executors for unsharded datasets are not managed by
+        SplitCoordinator actors and hence, are not accessible via the DatasetManager
+        so their cleanup is not handled by this method.
+        """
+        try:
+            shutdown_refs = [
+                coord.shutdown_executor.remote() for coord in self._coordinator_actors
+            ]
+            ray.get(shutdown_refs, timeout=5)
+        except GetTimeoutError:
+            logger.error("Ray Data executor shutdown task timed out after 5 seconds.")
+        except Exception:
+            logger.exception(
+                "Failed to gracefully terminate the Ray Data executor for each running dataset."
+            )

python/ray/train/v2/_internal/data_integration/interfaces.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Protocol, Union
+
+if TYPE_CHECKING:
+    from ray.data import DataIterator, Dataset
+
+# 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"]]
+
+
+@dataclass
+class DatasetShardMetadata:
+    """Metadata about a dataset shard used for lookup and configuration."""
+
+    dataset_name: str
+    world_rank: int
+
+
+class DatasetShardProvider(Protocol):
+    def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
+        """Get the dataset shard for the given dataset info.
+        Args:
+            dataset_info: The metadata of the shard to retrieve,
+                including the dataset name.
+        Returns:
+            The :class:`~ray.data.DataIterator` shard for the given dataset info.
+        Raises:
+            KeyError: If the dataset shard for the given dataset info is not found.
+        """
+        ...

python/ray/train/v2/_internal/exceptions.py

@@ -0,0 +1,165 @@
+import os
+from typing import List, Optional
+
+from ray.train.v2._internal.constants import (
+    COLLECTIVE_TIMEOUT_S_ENV_VAR,
+    DEFAULT_WORKER_GROUP_START_TIMEOUT_S,
+    DEFAULT_WORKER_HEALTH_CHECK_TIMEOUT_S,
+    WORKER_GROUP_START_TIMEOUT_S_ENV_VAR,
+    WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR,
+)
+
+
+# TODO: Distinguish between user and system exceptions.
+class RayTrainError(Exception):
+    """Base class for all Ray Train exceptions."""
+
+
+class WorkerHealthCheckTimeoutError(RayTrainError):
+    """Exception raised when a worker health check hangs for long enough."""
+
+    def __init__(self, message):
+        timeout = os.getenv(
+            WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR, DEFAULT_WORKER_HEALTH_CHECK_TIMEOUT_S
+        )
+        message += (
+            f"\nSet the {WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR} "
+            "environment variable to increase the timeout "
+            f"(current value: {timeout} seconds)."
+        )
+        super().__init__(message)
+
+
+class WorkerHealthCheckFailedError(RayTrainError):
+    """Exception raised when a worker health check fails."""
+
+    def __init__(self, message, failure: Exception):
+        super().__init__(message)
+        self._message = message
+        self.health_check_failure = failure
+
+    def __reduce__(self):
+        return (self.__class__, (self._message, self.health_check_failure))
+
+    def __str__(self):
+        return self._message + "\n" + str(self.health_check_failure)
+
+
+class WorkerGroupStartupTimeoutError(RayTrainError):
+    """Exception raised when the worker group startup times out.
+
+    Example scenario: 4 GPUs are detected in the cluster, but when the worker
+    are actually scheduled, one of the nodes goes down and only 3 GPUs are
+    available. One of the worker tasks may be stuck pending, until a timeout is reached.
+    """
+
+    def __init__(self, num_workers: int):
+        timeout = float(
+            os.environ.get(
+                WORKER_GROUP_START_TIMEOUT_S_ENV_VAR,
+                DEFAULT_WORKER_GROUP_START_TIMEOUT_S,
+            )
+        )
+        self.num_workers = num_workers
+        super().__init__(
+            f"The worker group startup timed out after {timeout} seconds waiting "
+            f"for {num_workers} workers. "
+            "Potential causes include: "
+            "(1) temporary insufficient cluster resources while waiting for "
+            "autoscaling (ignore this warning in this case), "
+            "(2) infeasible resource request where the provided `ScalingConfig` "
+            "cannot be satisfied), "
+            "and (3) transient network issues. "
+            f"Set the {WORKER_GROUP_START_TIMEOUT_S_ENV_VAR} "
+            "environment variable to increase the timeout."
+        )
+
+    def __reduce__(self):
+        return (self.__class__, (self.num_workers,))
+
+
+class WorkerGroupStartupFailedError(RayTrainError):
+    """Exception raised when the worker group fails to start.
+
+    Example scenario: A worker is scheduled onto a node that dies while
+    the worker actor is initializing.
+    """
+
+
+class InsufficientClusterResourcesError(RayTrainError):
+    """Exception raised when the cluster has insufficient resources.
+
+    Example scenario: A worker that requires 1 GPU is scheduled onto a cluster
+    that only has CPU worker node types.
+    """
+
+
+class CheckpointManagerInitializationError(RayTrainError):
+    """Exception raised when the checkpoint manager fails to initialize from a snapshot.
+
+    Example scenarios:
+    1. The checkpoint manager snapshot version is old and
+        incompatible with the current version of Ray Train.
+    2. The checkpoint manager snapshot JSON file is corrupted.
+    3. The checkpoint manager snapshot references checkpoints that cannot be found
+        in the run storage path.
+    """
+
+
+class CollectiveTimeoutError(RayTrainError):
+    """Exception raised when an internal Ray Train collective operation of
+    the worker group times out.
+    """
+
+
+class BroadcastCollectiveTimeoutError(CollectiveTimeoutError):
+    """Exception raised when the broadcast operation times out.
+
+    There are two main timeout examples:
+    1. If not all workers call `ray.train.report`, the entire worker group will
+        hang until the timeout before raising. This prevents indefinite worker
+        group hangs.
+    2. If a worker is slow in the training loop and fails to reach the broadcast
+        time, the collective will time out.
+    """
+
+    def __init__(
+        self, time_elapsed: Optional[float], missing_ranks: List[int], timeout_s: float
+    ):
+        self._time_elapsed = time_elapsed
+        self._missing_ranks = missing_ranks
+        self._timeout_s = timeout_s
+
+        message = (
+            f"The collective operation timed out after {time_elapsed:.2f} seconds. "
+            f"The following ranks have not joined the collective operation: {missing_ranks}\n"
+            f"You can set the timeout with the {COLLECTIVE_TIMEOUT_S_ENV_VAR} "
+            f"environment variable (current value: {timeout_s:.2f} seconds). "
+            "Disable the timeout by setting the environment variable to `None`."
+        )
+        super().__init__(message)
+
+    def __reduce__(self):
+        return (
+            self.__class__,
+            (self._time_elapsed, self._missing_ranks, self._timeout_s),
+        )
+
+
+class UserExceptionWithTraceback(RayTrainError):
+    """This class wraps a user code exception raised on the worker
+    with its original traceback string, for logging and debugging purposes.
+
+    This is needed because the original exception traceback is not serialized
+    with the exception when it is *returned* back to the main process.
+    """
+
+    def __init__(self, exc: BaseException, traceback_str: str):
+        self._base_exc = exc
+        self._traceback_str = traceback_str
+
+    def __reduce__(self):
+        return (self.__class__, (self._base_exc, self._traceback_str))
+
+    def __str__(self):
+        return self._traceback_str

python/ray/train/v2/_internal/execution/callback.py

@@ -0,0 +1,240 @@
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from ray.train.v2._internal.execution.training_report import _TrainingReport
+from ray.train.v2.api.callback import RayTrainCallback
+from ray.train.v2.api.config import ScalingConfig
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.execution.context import TrainRunContext
+    from ray.train.v2._internal.execution.controller import (
+        TrainControllerState,
+    )
+    from ray.train.v2._internal.execution.failure_handling import FailureDecision
+    from ray.train.v2._internal.execution.scaling_policy import ResizeDecision
+    from ray.train.v2._internal.execution.worker_group import (
+        ExecutionGroup,
+        ReplicaGroup,
+        Worker,
+        WorkerGroup,
+        WorkerGroupContext,
+        WorkerGroupPollStatus,
+    )
+    from ray.train.v2.api.result import Result
+
+
+@DeveloperAPI
+class ExecutionGroupCallback(RayTrainCallback):
+    """Base callback for execution groups (worker groups and replica groups)."""
+
+    def before_init_train_context(
+        self, workers: List["Worker"]
+    ) -> Dict[str, List[Any]]:
+        """Called before initializing the TrainContext for an execution group.
+
+        Return:
+            A dictionary of additional arguments for TrainContext.
+            The key is the argument name and the value is a list of argument values
+            to pass to the TrainContext constructor of each worker in the group.
+        """
+        return {}
+
+    def after_execution_group_start(self, execution_group: "ExecutionGroup"):
+        """Called after an execution group is started or replaced.
+        All workers in the execution group should be ready to execute tasks."""
+        pass
+
+    def before_execution_group_shutdown(self, execution_group: "ExecutionGroup"):
+        """Called before an execution group is shut down.
+        Workers may be dead at this point due to actor failures."""
+        pass
+
+
+@DeveloperAPI
+class WorkerGroupCallback(ExecutionGroupCallback):
+    @contextmanager
+    def on_worker_group_start(self):
+        yield
+
+    def before_worker_group_start(self, worker_group_context: "WorkerGroupContext"):
+        """Called before the worker group actors are initialized."""
+        pass
+
+    def after_worker_group_start(self, worker_group: "WorkerGroup"):
+        """Called after the worker group actors are initialized.
+        All workers should be ready to execute tasks."""
+        return self.after_execution_group_start(worker_group)
+
+    def after_worker_group_training_start(self, worker_group: "WorkerGroup"):
+        pass
+
+    @contextmanager
+    def on_worker_group_shutdown(self):
+        yield
+
+    def before_worker_group_shutdown(self, worker_group: "WorkerGroup"):
+        """Called before the worker group is shut down.
+        Workers may be dead at this point due to actor failures, so this method
+        should catch and handle exceptions if attempting to execute tasks."""
+        return self.before_execution_group_shutdown(worker_group)
+
+    def after_worker_group_shutdown(self, worker_group_context: "WorkerGroupContext"):
+        """Called after the worker group is shut down."""
+        pass
+
+    def after_worker_group_poll_status(
+        self, worker_group_status: "WorkerGroupPollStatus"
+    ):
+        pass
+
+    def before_worker_group_abort(self, worker_group_context: "WorkerGroupContext"):
+        """Called before the worker group is aborted."""
+        pass
+
+    def after_worker_group_abort(self, worker_group_context: "WorkerGroupContext"):
+        """Called after the worker group is aborted."""
+        pass
+
+
+@DeveloperAPI
+class ReplicaGroupCallback(ExecutionGroupCallback):
+    """Callback for replica group lifecycle events."""
+
+    def after_replica_group_start(self, replica_group: "ReplicaGroup"):
+        """Called after a replica group is started or replaced.
+        All workers in the replica group should be ready to execute tasks."""
+        return self.after_execution_group_start(replica_group)
+
+    def before_replica_group_shutdown(self, replica_group: "ReplicaGroup"):
+        """Called before a replica group is shut down.
+        Workers may be dead at this point due to actor failures."""
+        return self.before_execution_group_shutdown(replica_group)
+
+
+@DeveloperAPI
+class ControllerCallback(RayTrainCallback):
+    def after_controller_start(self, train_run_context: "TrainRunContext"):
+        """Called immediately after `TrainController.run` is called,
+        before the control loop starts executing."""
+        pass
+
+    # TODO(matthewdeng): Revisit this callback interface for better extensibility.
+    # This hook was added for the specific use case of setting a `label_selector`
+    # for new worker groups (e.g., for TPU reservations). The current interface is
+    # tightly coupled to this purpose and limits its reuse for other use-cases.
+    def on_controller_start_worker_group(
+        self, *, scaling_config: ScalingConfig, num_workers: int
+    ) -> Optional[Dict[str, str]]:
+        """Called by the TrainController before the worker group is started.
+
+        This hook can be used to perform setup that modifies the worker group's
+        placement, such as reserving an accelerator slice.
+
+        Args:
+            scaling_config: The scaling configuration for the run.
+            num_workers: The number of workers to be started.
+
+        Returns:
+            An optional dictionary defining a `label_selector`
+            to gang schedule the worker group on the reserved TPU slice.
+        """
+        return None
+
+    async def before_controller_shutdown(self):
+        """Called before `TrainController.run` exits,
+        after the control loop has exited."""
+        pass
+
+    def after_controller_state_update(
+        self,
+        previous_state: "TrainControllerState",
+        current_state: "TrainControllerState",
+    ):
+        """Called whenever the controller state is updated."""
+        pass
+
+    def before_controller_execute_failure_decision(
+        self,
+        failure_decision: "FailureDecision",
+    ):
+        """Called before the controller executes a failure decision."""
+        pass
+
+    def before_controller_execute_resize_decision(
+        self,
+        resize_decision: "ResizeDecision",
+    ):
+        """Called before the controller executes a resize decision."""
+        pass
+
+    def after_controller_finish(self, result: "Result"):
+        """Called after the training run completes, providing access to the final result.
+
+        Args:
+            result: The final training result containing metrics and checkpoint.
+        """
+        pass
+
+    def before_controller_abort(self):
+        """Called during `TrainController.abort` before the actor process exits."""
+        pass
+
+
+# TODO: consider consolidating all metrics into one dict, possibly with UDF
+@DeveloperAPI
+class ReportCallback(RayTrainCallback):
+    def after_report(
+        self,
+        training_report: _TrainingReport,
+        metrics: List[Dict[str, Any]],
+    ):
+        """Called after all workers have reported a training result.
+
+        Note that this differs from `after_worker_group_poll_status`,
+        which may only contain a subset of workers that have reported.
+        For example, if only rank 0 is performing checkpointing, then
+        rank 0 would report a training result the slowest.
+        """
+        pass
+
+
+@DeveloperAPI
+class WorkerCallback(RayTrainCallback):
+    """
+    Callbacks that are hooked to the worker event.
+
+    These callbacks are created on the train driver process and then
+    copied and passed to all the workers.
+    The execution of these callbacks happens on each of the workers,
+    not on the train driver process.
+    """
+
+    def after_init_train_context(self):
+        pass
+
+    def before_worker_shutdown(self):
+        pass
+
+
+@DeveloperAPI
+class TrainContextCallback(RayTrainCallback):
+    """
+    Callbacks that are hooked to the train context event.
+
+    These callbacks are created on the train driver process and then
+    copied and passed to all the workers.
+    The execution of these callbacks happens on the train context of the workers.
+    """
+
+    @contextmanager
+    def on_report(self):
+        yield
+
+    @contextmanager
+    def on_checkpoint_sync(self):
+        yield
+
+    @contextmanager
+    def on_checkpoint_transfer(self):
+        yield

python/ray/train/v2/_internal/execution/callback_manager.py

@@ -0,0 +1,67 @@
+import logging
+
+from ray.train.v2.api.exceptions import ControllerError
+
+logger = logging.getLogger(__name__)
+
+
+class CallbackManager:
+    def __init__(self, callbacks):
+        self._callbacks = callbacks
+
+    def _get_method(self, callback, hook_name: str):
+        """Look up a hook method on a callback, raising if missing."""
+        callback_name = type(callback).__name__
+        method = getattr(callback, hook_name, None)
+        if method is None or not callable(method):
+            raise ControllerError(
+                AttributeError(
+                    f"Callback '{callback_name}' hook '{hook_name}' is missing "
+                    "or not callable."
+                )
+            )
+        return method, callback_name
+
+    def invoke(self, hook_name: str, *args, **context) -> None:
+        for callback in self._callbacks:
+            method, callback_name = self._get_method(callback, hook_name)
+            try:
+                method(*args, **context)
+            except Exception as e:
+                # TODO: Enable configuration to suppress exceptions.
+                logger.exception(
+                    f"Exception raised in callback hook '{hook_name}' from callback "
+                    f"'{callback_name}'."
+                )
+                raise ControllerError(e) from e
+
+    async def async_invoke(self, hook_name: str, *args, **context) -> None:
+        for callback in self._callbacks:
+            method, callback_name = self._get_method(callback, hook_name)
+            try:
+                await method(*args, **context)
+            except Exception as e:
+                # TODO: Enable configuration to suppress exceptions.
+                logger.exception(
+                    f"Exception raised in callback hook '{hook_name}' from callback "
+                    f"'{callback_name}'."
+                )
+                raise ControllerError(e) from e
+
+    def invoke_best_effort(self, hook_name: str, *args, **context) -> None:
+        """Invoke a hook on every callback, logging and suppressing errors.
+
+        Unlike ``invoke``, this does not fail fast — every callback is
+        attempted even if earlier ones raise.  Used for cleanup hooks
+        (e.g. ``before_controller_abort``) where partial execution is
+        better than skipping remaining callbacks.
+        """
+        for callback in self._callbacks:
+            method, callback_name = self._get_method(callback, hook_name)
+            try:
+                method(*args, **context)
+            except Exception as e:
+                logger.exception(
+                    f"Error in callback hook '{hook_name}' from callback "
+                    f"'{callback_name}': {e}"
+                )

python/ray/train/v2/_internal/execution/checkpoint/checkpoint_manager.py

@@ -0,0 +1,656 @@
+import asyncio
+import json
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import ray
+from ray._common.pydantic_compat import BaseModel
+from ray._private.ray_constants import env_float
+from ray.air.config import CheckpointConfig
+from ray.train._checkpoint import Checkpoint
+from ray.train._internal.checkpoint_manager import (
+    _CheckpointManager,
+    _insert_into_sorted_list,
+)
+from ray.train._internal.session import _TrainingResult
+from ray.train.v2._internal.constants import (
+    COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
+    DEFAULT_COLLECTIVE_WARN_INTERVAL_S,
+)
+from ray.train.v2._internal.exceptions import CheckpointManagerInitializationError
+from ray.train.v2._internal.execution.callback import (
+    ReportCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.context import StorageContext
+from ray.train.v2._internal.execution.storage import _exists_at_fs_path, delete_fs_path
+from ray.train.v2._internal.execution.training_report import _TrainingReport
+from ray.train.v2._internal.execution.worker_group import Worker
+from ray.train.v2._internal.util import wait_with_logging
+from ray.train.v2.api.report_config import CheckpointConsistencyMode
+from ray.train.v2.api.reported_checkpoint import (
+    ReportedCheckpoint,
+    ReportedCheckpointStatus,
+)
+from ray.train.v2.api.validation_config import ValidationTaskConfig
+
+logger = logging.getLogger(__name__)
+
+
+GET_ALL_REPORTED_CHECKPOINTS_PERIODIC_WARNING = """
+`get_all_reported_checkpoints` has been waiting for all checkpoints to get to the {consistency_mode} state for {time_elapsed_s:.2f} s.
+You can set the {warn_interval_env_var} environment variable to change the frequency of this warning (current value: {warn_interval_s} s).
+"""
+
+
+class _TrainingResultState(BaseModel):
+    # Increment version if the schema changes
+    version: int = 0
+    checkpoint_dir_name: str
+    metrics: dict
+
+
+class _CheckpointManagerState(BaseModel):
+    ray_version: str = ray.__version__
+    checkpoint_results: List[_TrainingResultState]
+    checkpoint_report_indices: List[int]
+    latest_checkpoint_result: Optional[_TrainingResultState] = None
+    pending_training_results: List[_TrainingResultState]
+    pending_validation_specs: List[Union[bool, ValidationTaskConfig]]
+    current_report_index: int
+
+    # List of processed checkpoints based on if successfully validated,
+    #   timed out or failed due to an error or canceled for some reason.
+    validated_checkpoint_dir_names: List[str]
+    timed_out_validation_checkpoint_dir_names: List[str]
+    failed_validation_checkpoint_dir_names: List[str]
+
+
+def _get_training_result_from_state(
+    state: _TrainingResultState,
+    storage_context: StorageContext,
+) -> _TrainingResult:
+    """Get a TrainingResult object from a Pydantic state object."""
+    return _TrainingResult(
+        checkpoint=Checkpoint(
+            path=storage_context.build_checkpoint_path_from_name(
+                state.checkpoint_dir_name
+            ),
+            filesystem=storage_context.storage_filesystem,
+        ),
+        metrics=state.metrics,
+    )
+
+
+def _get_state_from_training_result(
+    training_result: _TrainingResult,
+    storage_context: StorageContext,
+) -> _TrainingResultState:
+    """Get a Pydantic state object from a TrainingResult object."""
+    return _TrainingResultState(
+        checkpoint_dir_name=storage_context.extract_checkpoint_dir_name_from_path(
+            training_result.checkpoint.path
+        ),
+        metrics=training_result.metrics,
+    )
+
+
+class CheckpointManager(_CheckpointManager, ReportCallback, WorkerGroupCallback):
+    def __init__(
+        self,
+        checkpoint_config: CheckpointConfig,
+        storage_context: StorageContext,
+    ):
+        self._storage_context = storage_context
+        self._checkpoint_config = checkpoint_config
+
+        # This tracks the number of report calls that have been processed
+        # for the current worker group.
+        self._current_report_index = 0
+
+        # Map from pending checkpoint to validation.
+        self._pending_training_results: Dict[
+            Checkpoint, Tuple[_TrainingResult, Union[bool, ValidationTaskConfig]]
+        ] = {}
+
+        # Set of checkpoints that have successfully completed, been timed out
+        #   or failed validation.
+        self._validated_checkpoints: set = set()
+        self._timed_out_validation_checkpoints: set = set()
+        self._failed_validation_checkpoints: set = set()
+
+        # Map from checkpoint to report index. Used to order checkpoints.
+        self._checkpoint_to_report_index = {}
+
+        self._condition = asyncio.Condition()
+
+        # Strong references to background tasks created via
+        # ``asyncio.create_task`` to prevent them from being garbage
+        # collected mid-execution. The event loop only keeps weak refs.
+        self._background_tasks: set = set()
+
+        self._collective_warn_interval_s = env_float(
+            COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
+            DEFAULT_COLLECTIVE_WARN_INTERVAL_S,
+        )
+
+        super().__init__(checkpoint_config)
+        # If the snapshot is found, the checkpoint manager will restore its state.
+        # TODO(xgui): CheckpointManager is used to save or restore the checkpoint manager state.
+        # We should sanity check if we should see old state in the storage folder.
+        self._maybe_load_state_from_storage()
+
+    def register_checkpoint(
+        self,
+        training_report: _TrainingReport,
+    ):
+        """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:
+            training_report: Training report to register.
+        """
+        checkpoint_result = _TrainingResult(
+            checkpoint=training_report.checkpoint,
+            metrics=training_report.metrics,
+        )
+        self._latest_checkpoint_result = checkpoint_result
+        self._checkpoint_to_report_index[
+            checkpoint_result.checkpoint
+        ] = self._current_report_index
+
+        if self._checkpoint_config.checkpoint_score_attribute is not None:
+            # 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,
+                checkpoint_to_report_index=self._checkpoint_to_report_index,
+            )
+        else:
+            # If no metric is provided, just append (ordering by time of registration).
+            self._checkpoint_results.append(checkpoint_result)
+
+        if training_report.validation:
+            self._pending_training_results[checkpoint_result.checkpoint] = (
+                checkpoint_result,
+                training_report.validation,
+            )
+
+        self._current_report_index += 1
+
+        self._save_state_and_delete_old_checkpoints()
+
+        self._notify()
+
+    def update_checkpoints_with_validation_result(
+        self,
+        checkpoint_updates: Dict[
+            Checkpoint, Tuple[Dict[str, Any], ReportedCheckpointStatus]
+        ],
+    ):
+        """Finalize pending validations by recording terminal status and metrics.
+
+        * For VALIDATED checkpoints, metrics are merged into the checkpoint's
+          existing metrics and the checkpoint is re-sorted.
+        * For VALIDATION_TIMEOUT and VALIDATION_FAILED checkpoints, metrics are
+          left untouched and the checkpoint retains its original training-time
+          score position.
+        """
+        for checkpoint, (metrics, status) in checkpoint_updates.items():
+            if checkpoint not in self._pending_training_results:
+                logger.warning(
+                    f"Checkpoint {checkpoint} not found in pending training results. "
+                )
+                continue
+            checkpoint_result, _ = self._pending_training_results[checkpoint]
+            if checkpoint_result not in self._checkpoint_results:
+                raise ValueError(
+                    f"Checkpoint {checkpoint} was in pending training results but not "
+                    "checkpoint results. "
+                )
+            self._pending_training_results.pop(checkpoint)
+
+            if status == ReportedCheckpointStatus.VALIDATED:
+                # Update the metrics and sort into checkpoint_results
+                checkpoint_result.metrics.update(metrics)
+                self._checkpoint_results.remove(checkpoint_result)
+                _insert_into_sorted_list(
+                    self._checkpoint_results,
+                    checkpoint_result,
+                    key=self._get_checkpoint_score,
+                    checkpoint_to_report_index=self._checkpoint_to_report_index,
+                )
+                self._validated_checkpoints.add(checkpoint)
+            elif status == ReportedCheckpointStatus.VALIDATION_TIMEOUT:
+                self._timed_out_validation_checkpoints.add(checkpoint)
+            elif status == ReportedCheckpointStatus.VALIDATION_FAILED:
+                self._failed_validation_checkpoints.add(checkpoint)
+            else:
+                raise ValueError(
+                    f"Unexpected terminal validation status {status} for "
+                    f"checkpoint {checkpoint}."
+                )
+
+        self._save_state_and_delete_old_checkpoints()
+        self._notify()
+
+    def get_pending_training_results(
+        self,
+    ) -> Dict[Checkpoint, Tuple[_TrainingResult, Union[bool, ValidationTaskConfig]]]:
+        """Get the pending training results which includes their validation specs."""
+        return self._pending_training_results
+
+    def _notify(self):
+        """Notify condition so all listeners know state has changed."""
+
+        async def async_notify():
+            async with self._condition:
+                self._condition.notify_all()
+
+        # Keep a strong reference to the task so it isn't garbage
+        # collected before completing, which would silently drop
+        # the notification and could leave listeners waiting forever.
+        task = asyncio.create_task(async_notify())
+        self._background_tasks.add(task)
+        task.add_done_callback(self._background_tasks.discard)
+
+    def _save_state_and_delete_old_checkpoints(self):
+        """Delete the old checkpoints."""
+        # Get checkpoints to delete
+        results_to_delete = set()
+        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 and pending checkpoints
+            results_to_delete = worst_results - {self._latest_checkpoint_result}
+            results_to_delete = results_to_delete - {
+                v for v, _ in self._pending_training_results.values()
+            }
+
+            # 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:
+                del self._checkpoint_to_report_index[checkpoint_result.checkpoint]
+
+                # discard doesn't raise an error if the element isn't found
+                self._validated_checkpoints.discard(checkpoint_result.checkpoint)
+                self._timed_out_validation_checkpoints.discard(
+                    checkpoint_result.checkpoint
+                )
+                self._failed_validation_checkpoints.discard(
+                    checkpoint_result.checkpoint
+                )
+
+        # Save the checkpoint manager state to storage.
+        # Note: We save the state before deleting the old checkpoints.
+        # If deletion happens first and the process crashes, our snapshot
+        # may point to some stale checkpoints that are already deleted.
+        # TODO: Make this writing operation non-blocking.
+        self._write_state_to_storage()
+
+        # Delete the old checkpoints.
+        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)
+
+    # --------------------------
+    # CheckpointManager state
+    # --------------------------
+
+    def _save_state(self) -> str:
+        """Save the checkpoint manager state to a JSON str."""
+
+        checkpoint_results = [
+            _get_state_from_training_result(checkpoint_result, self._storage_context)
+            for checkpoint_result in self._checkpoint_results
+        ]
+
+        checkpoint_report_indices = [
+            self._checkpoint_to_report_index[checkpoint_result.checkpoint]
+            for checkpoint_result in self._checkpoint_results
+        ]
+
+        latest_checkpoint_result = (
+            _get_state_from_training_result(
+                self._latest_checkpoint_result, self._storage_context
+            )
+            if self._latest_checkpoint_result is not None
+            else None
+        )
+
+        pending_training_results = [
+            _get_state_from_training_result(v, self._storage_context)
+            for v, _ in self._pending_training_results.values()
+        ]
+        pending_validation_specs = [
+            v for _, v in self._pending_training_results.values()
+        ]
+
+        validated_ckpt_dir_names = [
+            self._storage_context.extract_checkpoint_dir_name_from_path(checkpoint.path)
+            for checkpoint in self._validated_checkpoints
+        ]
+        timed_out_validation_ckpt_dir_names = [
+            self._storage_context.extract_checkpoint_dir_name_from_path(checkpoint.path)
+            for checkpoint in self._timed_out_validation_checkpoints
+        ]
+        failed_validation_ckpt_dir_names = [
+            self._storage_context.extract_checkpoint_dir_name_from_path(checkpoint.path)
+            for checkpoint in self._failed_validation_checkpoints
+        ]
+
+        manager_snapshot = _CheckpointManagerState(
+            checkpoint_results=checkpoint_results,
+            checkpoint_report_indices=checkpoint_report_indices,
+            latest_checkpoint_result=latest_checkpoint_result,
+            pending_training_results=pending_training_results,
+            pending_validation_specs=pending_validation_specs,
+            current_report_index=self._current_report_index,
+            validated_checkpoint_dir_names=validated_ckpt_dir_names,
+            timed_out_validation_checkpoint_dir_names=timed_out_validation_ckpt_dir_names,
+            failed_validation_checkpoint_dir_names=failed_validation_ckpt_dir_names,
+        )
+        return manager_snapshot.json()
+
+    def _load_state(self, json_state: str):
+        """Load the checkpoint manager state from a JSON str."""
+        json_dict = None
+        try:
+            json_dict = json.loads(json_state)
+            manager_snapshot = _CheckpointManagerState.parse_obj(json_dict)
+        except Exception as e:
+            if not json_dict:
+                error = e
+            elif "ray_version" not in json_dict:
+                error = (
+                    "You are loading a checkpoint manager snapshot saved with an unknown Ray version "
+                    f"but you are running Ray version {ray.__version__}. Please use the same Ray version "
+                    "the checkpoint manager snapshot was saved with."
+                )
+            elif json_dict["ray_version"] != ray.__version__:
+                error = (
+                    f"You are loading a checkpoint manager snapshot saved with Ray version "
+                    f"{json_dict['ray_version']} but you are running Ray version "
+                    f"{ray.__version__}. Please use the same Ray version the checkpoint "
+                    "manager snapshot was saved with."
+                )
+            else:
+                error = e
+            raise CheckpointManagerInitializationError(error) from e
+
+        # Do this so we are using the same checkpoint and trainingresult objects.
+        # TODO: consider asserting that every checkpoint has a unique dir name
+        checkpoint_dir_name_to_checkpoint_result = {}
+
+        for training_result_state in manager_snapshot.checkpoint_results:
+            training_result = _get_training_result_from_state(
+                training_result_state, self._storage_context
+            )
+            checkpoint_dir_name_to_checkpoint_result[
+                training_result_state.checkpoint_dir_name
+            ] = training_result
+            self._checkpoint_results.append(training_result)
+        self._assert_checkpoints_exist()
+
+        assert len(self._checkpoint_results) == len(
+            manager_snapshot.checkpoint_report_indices
+        )
+        self._checkpoint_to_report_index = {
+            checkpoint_result.checkpoint: report_index
+            for checkpoint_result, report_index in zip(
+                self._checkpoint_results, manager_snapshot.checkpoint_report_indices
+            )
+        }
+
+        self._latest_checkpoint_result = (
+            checkpoint_dir_name_to_checkpoint_result[
+                manager_snapshot.latest_checkpoint_result.checkpoint_dir_name
+            ]
+            if manager_snapshot.latest_checkpoint_result is not None
+            else None
+        )
+
+        assert len(manager_snapshot.pending_training_results) == len(
+            manager_snapshot.pending_validation_specs
+        )
+        for training_result_state, validation_spec in zip(
+            manager_snapshot.pending_training_results,
+            manager_snapshot.pending_validation_specs,
+        ):
+            training_result = checkpoint_dir_name_to_checkpoint_result[
+                training_result_state.checkpoint_dir_name
+            ]
+            self._pending_training_results[training_result.checkpoint] = (
+                training_result,
+                validation_spec,
+            )
+
+        # Restore terminal validation statuses. Only checkpoints still in
+        # _checkpoint_results can be looked up; evicted checkpoints are irrelevant.
+        for dir_names, target_set in (
+            (
+                manager_snapshot.validated_checkpoint_dir_names,
+                self._validated_checkpoints,
+            ),
+            (
+                manager_snapshot.timed_out_validation_checkpoint_dir_names,
+                self._timed_out_validation_checkpoints,
+            ),
+            (
+                manager_snapshot.failed_validation_checkpoint_dir_names,
+                self._failed_validation_checkpoints,
+            ),
+        ):
+            for dir_name in dir_names:
+                if dir_name in checkpoint_dir_name_to_checkpoint_result:
+                    target_set.add(
+                        checkpoint_dir_name_to_checkpoint_result[dir_name].checkpoint
+                    )
+
+        self._current_report_index = manager_snapshot.current_report_index
+
+    def _maybe_load_state_from_storage(self):
+        """Load the checkpoint manager state from storage.
+        If no snapshot is found, start with a clean state.
+        """
+        if not _exists_at_fs_path(
+            fs=self._storage_context.storage_filesystem,
+            fs_path=self._storage_context.checkpoint_manager_snapshot_path,
+        ):
+            logger.debug(
+                "No checkpoint manager snapshot found. "
+                "No checkpoint will be available via `ray.train.get_checkpoint`, "
+                "so training will start from scratch."
+            )
+            return
+        with self._storage_context.storage_filesystem.open_input_stream(
+            self._storage_context.checkpoint_manager_snapshot_path
+        ) as f:
+            logger.info(
+                "A run snapshot was found in storage folder at: "
+                f"'{self._storage_context.experiment_fs_path}'\n"
+                "This snapshot contains a list of checkpoints reported via "
+                "`ray.train.report` and will be loaded. "
+                "This allows the latest checkpoint found in the snapshot to be "
+                "accessible within your training function via "
+                "`ray.train.get_checkpoint`.\n"
+                "If you meant to start a brand new training job without any "
+                "information about previous checkpoints found in this directory, "
+                "please configure a new, unique `RunConfig(name)` or delete the "
+                f"existing folder at '{self._storage_context.experiment_fs_path}'."
+            )
+            json_state = f.read().decode("utf-8")
+        self._load_state(json_state)
+
+    def _write_state_to_storage(self):
+        """Write the checkpoint manager state to storage."""
+        checkpoint_manager_snapshot = self._save_state()
+        with self._storage_context.storage_filesystem.open_output_stream(
+            self._storage_context.checkpoint_manager_snapshot_path
+        ) as f:
+            f.write(checkpoint_manager_snapshot.encode("utf-8"))
+
+    def _assert_checkpoints_exist(self):
+        """Validate the checkpoint manager state.
+
+        This method will validate the checkpoint manager state by checking if
+        the checkpoints specified in manager snapshot is compatible with the
+        checkpoint folders of the experiment storage filesystem.
+
+        Raises:
+            CheckpointManagerInitializationError: If the checkpoint manager snapshot
+                is not consistent with the stored checkpoints.
+        """
+        for checkpoint_result in self._checkpoint_results:
+            checkpoint = checkpoint_result.checkpoint
+            assert checkpoint is not None
+            if not _exists_at_fs_path(
+                fs=checkpoint.filesystem, fs_path=checkpoint.path
+            ):
+                raise CheckpointManagerInitializationError(
+                    "The run snapshot contains a reference to a checkpoint "
+                    f"that does not exist anymore ({checkpoint}). You are "
+                    "running in a corrupted run directory `experiment_fs_path`. "
+                    "Please configure a new, unique `RunConfig(name)` "
+                    "or delete the existing folder at "
+                    f"`{self._storage_context.experiment_fs_path}`."
+                )
+
+    # --------------------------
+    # ReportCallback
+    # --------------------------
+
+    def after_report(
+        self,
+        training_report: _TrainingReport,
+        metrics: List[Dict[str, Any]],
+    ):
+        if not training_report.checkpoint:
+            self._current_report_index += 1
+            self._notify()
+            return
+
+        self.register_checkpoint(training_report)
+
+    # --------------------------
+    # WorkerGroupCallback
+    # --------------------------
+
+    def before_init_train_context(self, workers: List[Worker]) -> Dict[str, List[Any]]:
+        latest_checkpoint = (
+            self.latest_checkpoint_result.checkpoint
+            if self.latest_checkpoint_result
+            else None
+        )
+        train_context_args = {
+            "checkpoint": [latest_checkpoint] * len(workers),
+            "current_report_index": [self._current_report_index] * len(workers),
+        }
+        return train_context_args
+
+    # --------------------------------
+    # Get all reported checkpoints API
+    # --------------------------------
+
+    def _get_checkpoint_status(
+        self, checkpoint: Checkpoint
+    ) -> ReportedCheckpointStatus:
+        """Get ReportedCheckpoint's status."""
+        if checkpoint in self._pending_training_results:
+            return ReportedCheckpointStatus.PENDING_VALIDATION
+        elif checkpoint in self._timed_out_validation_checkpoints:
+            return ReportedCheckpointStatus.VALIDATION_TIMEOUT
+        elif checkpoint in self._failed_validation_checkpoints:
+            return ReportedCheckpointStatus.VALIDATION_FAILED
+        elif checkpoint in self._validated_checkpoints:
+            return ReportedCheckpointStatus.VALIDATED
+        else:
+            return ReportedCheckpointStatus.COMMITTED
+
+    def _generate_get_all_reported_checkpoints_periodic_warning(
+        self, start_time: float, consistency_mode: CheckpointConsistencyMode
+    ) -> str:
+        """Generates the warning message for the get_all_reported_checkpoints periodic warning."""
+        return GET_ALL_REPORTED_CHECKPOINTS_PERIODIC_WARNING.format(
+            consistency_mode=consistency_mode,
+            time_elapsed_s=asyncio.get_event_loop().time() - start_time,
+            warn_interval_env_var=COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
+            warn_interval_s=self._collective_warn_interval_s,
+        )
+
+    async def get_all_reported_checkpoints(
+        self,
+        current_report_index: int,
+        consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
+        timeout_s: Optional[float] = None,
+    ) -> List[ReportedCheckpoint]:
+        """Get all the reported checkpoints so far.
+
+        Args:
+            current_report_index: The current report index.
+            consistency_mode: Read semantics for checkpoint retrieval. Defaults to VALIDATED.
+            timeout_s: Timeout in seconds. Defaults to None to run forever.
+
+        Returns:
+            A list of ReportedCheckpoint objects that represent the checkpoints and
+            corresponding metrics reported by the workers.
+        """
+        start_time = asyncio.get_event_loop().time()
+        if consistency_mode == CheckpointConsistencyMode.COMMITTED:
+
+            def predicate() -> bool:
+                return self._current_report_index == current_report_index
+
+        elif consistency_mode == CheckpointConsistencyMode.VALIDATED:
+
+            def predicate() -> bool:
+                return (
+                    self._current_report_index == current_report_index
+                    and not self._pending_training_results
+                )
+
+        else:
+            raise ValueError(
+                f"Unexpected CheckpointConsistencyMode: {consistency_mode}"
+            )
+
+        async with self._condition:
+            try:
+                await wait_with_logging(
+                    self._condition,
+                    predicate=predicate,
+                    generate_warning_message=lambda: self._generate_get_all_reported_checkpoints_periodic_warning(
+                        start_time, consistency_mode
+                    ),
+                    warn_interval_s=self._collective_warn_interval_s,
+                    timeout_s=timeout_s,
+                )
+            except (asyncio.TimeoutError, TimeoutError):
+                # Time out due to checkpoint upload or validation in progress
+                logger.debug(
+                    "Timed out waiting for reported_checkpoint to become available."
+                )
+
+        # TODO: might be nice for CheckpointManager to manage ReportedCheckpoint
+        # instead of _TrainingResult but that is a large refactor.
+        return [
+            ReportedCheckpoint(
+                checkpoint=tr.checkpoint,
+                metrics=tr.metrics,
+                status=self._get_checkpoint_status(tr.checkpoint),
+            )
+            for tr in self._checkpoint_results
+        ]

python/ray/train/v2/_internal/execution/checkpoint/report_handler.py

@@ -0,0 +1,129 @@
+from collections import deque
+from typing import Deque, List, Optional
+
+from ray.train.v2._internal.execution.callback import (
+    ReplicaGroupCallback,
+    ReportCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.training_report import _TrainingReport
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroup,
+    WorkerGroupPollStatus,
+)
+from ray.train.v2._internal.execution.worker_group.execution_group import ReplicaGroup
+
+
+class ReportCallbackHandler(ReplicaGroupCallback, WorkerGroupCallback):
+    """Consolidate training results from multiple workers and call
+    subscribers implementing the `ReportCallback` interface sequentially.
+    """
+
+    def __init__(self, report_callbacks: List[ReportCallback]):
+        # We set the worker group after it has been started and remove it after it
+        # has been shut down.
+        self._worker_group: Optional[WorkerGroup] = None
+        # A list of queues holding training reports from workers.
+        self._training_report_queues: Optional[List[Deque[_TrainingReport]]] = None
+
+        self._report_callbacks = report_callbacks
+
+    def _assert_initialized(self):
+        assert (
+            self._worker_group and self._training_report_queues
+        ), "Need to call initialize state with `after_worker_group_start` first."
+
+    # --------------------------
+    # WorkerGroupCallback
+    # --------------------------
+
+    def after_worker_group_poll_status(
+        self, worker_group_status: WorkerGroupPollStatus
+    ) -> None:
+        """Handle training results as they roll in from worker status polls.
+
+        Wait for all workers to report training results to collect
+        a consolidated training result.
+        """
+        # Step 1: Assert that the worker group has been started and not shut down.
+        self._assert_initialized()
+
+        assert len(self._worker_group) == len(worker_group_status.worker_statuses), (
+            f"The number of workers in the worker group has changed unexpectedly. "
+            f"Expected: {len(self._worker_group)}, got: {len(worker_group_status.worker_statuses)}"
+        )
+
+        # Step 2: Update training_reports_queues with poll_results.
+        for i in range(len(self._worker_group)):
+            training_report = worker_group_status.worker_statuses[i].training_report
+            if training_report:
+                self._training_report_queues[i].append(training_report)
+
+        # Directly return if any of the worker result queues are empty.
+        if not all(self._training_report_queues):
+            return
+
+        training_reports = [q.popleft() for q in self._training_report_queues]
+
+        # Step 3: Consolidate a list of checkpoints to single checkpoint.
+        # Use the first checkpoint as the consolidated checkpoint.
+        checkpoint_results = [
+            tr for tr in training_reports if tr.checkpoint is not None
+        ]
+
+        consolidated_checkpoint = None
+        validation = False
+        if checkpoint_results:
+            # Double check the storage path of the checkpoints in the training results.
+            unique_checkpoint_paths = {tr.checkpoint.path for tr in checkpoint_results}
+            if len(unique_checkpoint_paths) > 1:
+                # TODO: Support for inconsistent checkpoints path from workers
+                # instead of hard raising error. Maybe drop this iteration of
+                # training results and continue with the next iteration.
+                raise RuntimeError(
+                    "The storage path of the checkpoints in the training results "
+                    "is not the same. This means the checkpoints are not consistent."
+                    "Got a mix of the following checkpoint paths: "
+                    f"{unique_checkpoint_paths}\n"
+                    "This is unexpected -- please file a Github issue."
+                )
+            consolidated_checkpoint = checkpoint_results[0].checkpoint
+            validation = checkpoint_results[0].validation
+
+        # Step 4: Invoke all dependent `ReportCallback`s.
+        metrics_per_worker = [
+            training_report.metrics for training_report in training_reports
+        ]
+        for callback in self._report_callbacks:
+            callback.after_report(
+                training_report=_TrainingReport(
+                    checkpoint=consolidated_checkpoint,
+                    metrics=metrics_per_worker[0],
+                    validation=validation,
+                ),
+                metrics=metrics_per_worker,
+            )
+
+    def after_worker_group_start(self, worker_group: WorkerGroup) -> None:
+        """Handle worker group start. Initialize internal states."""
+        self._worker_group = worker_group
+        self._training_report_queues = [deque() for _ in range(len(self._worker_group))]
+
+    def before_worker_group_shutdown(self, worker_group: WorkerGroup) -> None:
+        """Handle worker group shutdown. Clear internal states.
+
+        None of the partial reported results are valid at this point.
+        """
+        self._worker_group = None
+        self._training_report_queues = None
+
+    # --------------------------
+    # ReplicaGroupCallback
+    # --------------------------
+
+    def after_replica_group_start(self, replica_group: ReplicaGroup) -> None:
+        """Handle replica group start. Initialize internal states."""
+        self._assert_initialized()
+        # TODO: it might be possible to reuse existing queues.
+        # For example, if 3/4 ddp workers reported a checkpoint, that checkpoint is usable.
+        self._training_report_queues = [deque() for _ in range(len(self._worker_group))]

python/ray/train/v2/_internal/execution/checkpoint/sync_actor.py

@@ -0,0 +1,226 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from typing import List, Optional, TypeVar
+
+import ray
+from ray.train.v2._internal.constants import (
+    COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
+    DEFAULT_COLLECTIVE_TIMEOUT_S,
+    DEFAULT_COLLECTIVE_WARN_INTERVAL_S,
+)
+from ray.train.v2._internal.exceptions import BroadcastCollectiveTimeoutError
+from ray.train.v2._internal.util import wait_with_logging
+
+T = TypeVar("T", bound=Optional[object])
+logger = logging.getLogger(__name__)
+
+
+class SynchronizationBarrierResetError(Exception):
+    """Raised when the synchronization barrier is reset, e.g. due to a worker failure."""
+
+    pass
+
+
+BROADCAST_PERIODIC_WARNING = """
+`{caller_method_name}` has not been called by all {world_size} workers in the group.
+The workers have been waiting for {max_time_elapsed_s:.2f} s for the following ranks to join the `{caller_method_name}` call: {missing_ranks}.
+Also ensure that workers are not hanging on other operations, causing them to miss this synchronization barrier.
+You can set the {warn_interval_env_var} environment variable to change the frequency of this warning (current value: {warn_interval_s} s).
+"""
+
+
+@ray.remote(num_cpus=0)  # type: ignore
+class SynchronizationActor:
+    """A Ray actor that synchronizes the workers in a distributed training job.
+
+    This actor forms a synchronization barrier on a group of processes.
+    Every time a worker calls the broadcast_from_rank_zero method,
+    the counter is incremented. When the counter equals to the world size,
+    the actor notifies all the workers to continue.
+    """
+
+    def __init__(
+        self,
+        timeout_s: Optional[float] = DEFAULT_COLLECTIVE_TIMEOUT_S,
+        warn_interval_s: float = DEFAULT_COLLECTIVE_WARN_INTERVAL_S,
+    ):
+        self._counter: int = 0
+        self._world_size: int = 0
+        self._condition = asyncio.Condition()
+        self._reduced_data = None
+        self._reset = False
+        # The time when workers from different ranks
+        # enters the synchronization barrier.
+        self._sync_start_times: List[Optional[float]] = []
+        # The timeout in seconds for the synchronization barrier.
+        self._timeout_s: Optional[float] = timeout_s
+        # The interval in seconds to log a warning when waiting for the barrier.
+        self._warn_interval_s: float = warn_interval_s
+
+    def get_counter(self):
+        """Returns the current value of the counter."""
+        return self._counter
+
+    def get_world_size(self):
+        """Returns the current value of the world_size."""
+        return self._world_size
+
+    def get_reduced_data(self):
+        """Returns the current value of the reduced_data."""
+        return self._reduced_data
+
+    def _clear_states(self):
+        """Clears the states of the actor. When the last worker has
+        called the _clear_states method, the actor clears its states
+        """
+        self._counter -= 1
+        if self._counter == 0:
+            self._reduced_data = None
+            self._world_size = 0
+            self._reset = False
+            self._condition.notify_all()
+
+    async def _setup_or_validate_collective_op(self, world_size: int):
+        """The setup method for the synchronization actor if it is not setup yet.
+        It initializes the world size and the start times for the
+        synchronization barrier.
+        """
+        # Wait for previous collective reset to finish.
+        await self._condition.wait_for(lambda: not self._reset)
+        if self._world_size == 0:
+            self._world_size = world_size
+            self._sync_start_times = [None] * world_size
+        elif world_size != self._world_size:
+            raise ValueError(
+                f"Expected all callers to provide the same world size. \
+                Got {world_size} and expected {self._world_size}."
+            )
+
+    @asynccontextmanager
+    async def _broadcast_collective_context_manager(
+        self, world_rank: int, world_size: int, data: T
+    ):
+        """A context manager that ensures the synchronization barrier is lifted
+        after the block of code is executed.
+        """
+        try:
+            await self._setup_or_validate_collective_op(world_size)
+            if world_rank == 0:
+                self._reduced_data = data
+            if self._counter < self._world_size:
+                self._counter += 1
+            yield
+        finally:
+            self._clear_states()
+
+    def _get_time_elapsed(self) -> Optional[float]:
+        """Return the time elapsed since the first worker entered the barrier.
+        If no workers have entered the barrier, returns None.
+        """
+        start_times = [t for t in self._sync_start_times if t is not None]
+        if not start_times:
+            return None
+
+        return asyncio.get_event_loop().time() - min(start_times)
+
+    def _get_missing_ranks(self) -> List[int]:
+        """Returns the ranks that have not entered the synchronization barrier."""
+        return [i for i, t in enumerate(self._sync_start_times) if t is None]
+
+    def _generate_broadcast_periodic_warning(self, caller_method_name: str) -> str:
+        """Generates the warning message for the broadcast periodic warning."""
+
+        return BROADCAST_PERIODIC_WARNING.format(
+            caller_method_name=caller_method_name,
+            world_size=self._world_size,
+            max_time_elapsed_s=self._get_time_elapsed(),
+            missing_ranks=self._get_missing_ranks(),
+            warn_interval_env_var=COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
+            warn_interval_s=self._warn_interval_s,
+        )
+
+    async def reset(self):
+        """Reset the synchronization barrier, unblocking any waiting workers.
+
+        If no workers are currently at the barrier, this is a no-op.
+        Waiting workers will raise SynchronizationBarrierResetError.
+        The actor remains alive and usable for subsequent barriers.
+        """
+        async with self._condition:
+            if self._counter == 0:
+                return
+            self._reset = True
+            self._condition.notify_all()
+
+    async def broadcast_from_rank_zero(
+        self,
+        world_rank: int,
+        world_size: int,
+        data: T,
+        caller_method_name: str,
+    ) -> T:
+        """Broadcasts a data from the worker with rank 0 to all other workers.
+
+        This method is a coroutine that blocks until all workers have called this
+        method  with the their data. The data from the worker with rank 0 will
+        be returned.
+
+        Args:
+            world_rank: The rank of the worker that calls this method.
+            world_size: The total number of workers in the group.
+            data: The data to broadcast.
+            caller_method_name: The name of the method that calls this method.
+
+        Returns:
+            The data broadcasted from the worker with rank 0.
+        """
+        # TODO: resolve https://github.com/ray-project/ray/pull/54066#discussion_r2180657435
+        # We couldn't reproduce the issue but the asyncio docs don't say it can't happen.
+
+        # Ensures that all global states manipulation is done within the async context
+        # manager which makes the condition variable awaiting and the counter
+        # incrementing an atomic operation.
+        async with self._condition:
+            async with self._broadcast_collective_context_manager(
+                world_rank, world_size, data
+            ):
+                # If the counter is equal to the world size, it means the last worker
+                # has called the broadcast_from_rank_zero method. The actor notifies
+                # all the workers to continue.
+                if self._counter == self._world_size:
+                    self._condition.notify_all()
+                    return self._reduced_data
+                # If the counter is less than the world size, the actor waits for the
+                # other workers to call the broadcast_from_rank_zero method.
+                try:
+                    current_time = asyncio.get_event_loop().time()
+                    self._sync_start_times[world_rank] = current_time
+                    await wait_with_logging(
+                        self._condition,
+                        predicate=None,
+                        generate_warning_message=(
+                            lambda: self._generate_broadcast_periodic_warning(
+                                caller_method_name
+                            )
+                        )
+                        if world_rank == 0
+                        else None,
+                        warn_interval_s=self._warn_interval_s,
+                        timeout_s=self._timeout_s,
+                    )
+                    if self._reset:
+                        raise SynchronizationBarrierResetError(
+                            "Synchronization barrier was reset, likely due "
+                            "to a worker failure and replica group replacement."
+                        )
+                    return self._reduced_data
+                except (asyncio.TimeoutError, TimeoutError) as e:
+                    raise BroadcastCollectiveTimeoutError(
+                        time_elapsed=self._get_time_elapsed(),
+                        missing_ranks=self._get_missing_ranks(),
+                        timeout_s=self._timeout_s,
+                    ) from e
+
+    # TODO: Implement a general consensus_from_votes method that takes a callable
+    # reduce_fn and a list of votes from each worker. The method returns the consensus

python/ray/train/v2/_internal/execution/checkpoint/validation_manager.py

@@ -0,0 +1,289 @@
+import asyncio
+import logging
+import time
+from collections import OrderedDict, deque
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import ray
+from ray.train._checkpoint import Checkpoint
+from ray.train.v2._internal.execution.callback import (
+    ControllerCallback,
+    ReportCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
+    CheckpointManager,
+)
+from ray.train.v2._internal.execution.training_report import (
+    _TrainingReport,
+)
+from ray.train.v2.api.reported_checkpoint import ReportedCheckpointStatus
+from ray.train.v2.api.validation_config import ValidationConfig, ValidationTaskConfig
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.execution.controller import TrainControllerState
+    from ray.train.v2._internal.execution.worker_group.worker import Worker
+
+logger = logging.getLogger(__name__)
+
+
+VALIDATION_TASK_POLL_INTERVAL_S = 1
+MAX_IN_FLIGHT_VALIDATIONS = 1
+
+
+@dataclass
+class _PendingValidation:
+    checkpoint: Checkpoint
+    start_time: float
+    # None when no timeout applies.
+    timeout_s: Optional[float]
+
+    def __post_init__(self):
+        assert (
+            self.timeout_s is None or self.timeout_s > 0
+        ), f"timeout_s needs to be None (for no timeout) or a positive value in seconds. Actual value: {self.timeout_s}"
+
+
+@ray.remote
+def run_validation_fn(
+    validation_config: ValidationConfig,
+    validation_task_config: Union[bool, ValidationTaskConfig],
+    checkpoint: Checkpoint,
+) -> Dict:
+    """Run the user-defined validation function.
+
+    Merges fn_kwargs from validation_config.task_config (defaults) with
+    fn_kwargs from validation_task_config (per-report overrides).
+    """
+    # Merge kwargs: defaults from validation_config, overrides from validation_task_config
+    if validation_task_config is True:
+        merged_kwargs = validation_config.task_config.fn_kwargs
+    else:
+        merged_kwargs = {
+            **validation_config.task_config.fn_kwargs,
+            **validation_task_config.fn_kwargs,
+        }
+    metrics_dict = validation_config.fn(
+        checkpoint,
+        **merged_kwargs,
+    )
+    if not isinstance(metrics_dict, dict):
+        raise ValueError(
+            "The validation function must return a dictionary of metrics. "
+            f"Got {type(metrics_dict)} instead."
+        )
+    return metrics_dict
+
+
+class ValidationManager(ControllerCallback, ReportCallback, WorkerGroupCallback):
+    def __init__(
+        self,
+        checkpoint_manager: CheckpointManager,
+        validation_config: ValidationConfig,
+    ):
+        self._checkpoint_manager = checkpoint_manager
+        self._validation_config = validation_config
+
+        # _TrainingReports that we will validate
+        self._training_report_queue = deque()
+
+        # Map from in flight validation task to its pending-validation record
+        # (checkpoint + start_time + resolved timeout).
+        self._pending_validations: "OrderedDict[ray.ObjectRef, _PendingValidation]" = (
+            OrderedDict()
+        )
+
+        # Tasks that this manager proactively cancelled due to timeout. Used to
+        # distinguish timeout-cancels from controller-abort-cancels (both raise
+        # TaskCancelledError on ray.get).
+        self._timed_out_tasks: set = set()
+
+        # Map from validation task to checkpoint
+        # Finished validations that have yet to be processed
+        self._finished_validations: "OrderedDict[ray.ObjectRef, Checkpoint]" = (
+            OrderedDict()
+        )
+
+        self._requeue_incomplete_validations()
+
+    def _requeue_incomplete_validations(self):
+        """Add _TrainingReports for incomplete validations to the queue."""
+        for checkpoint, (
+            training_result,
+            validation,
+        ) in self._checkpoint_manager.get_pending_training_results().items():
+            if validation:
+                self._training_report_queue.append(
+                    _TrainingReport(
+                        metrics=training_result.metrics,
+                        checkpoint=checkpoint,
+                        validation=validation,
+                    )
+                )
+
+    def after_report(
+        self,
+        training_report: _TrainingReport,
+        metrics: List[Dict[str, Any]],
+    ):
+        if training_report.validation:
+            self._training_report_queue.append(training_report)
+
+    def _cancel_timed_out_validations(self):
+        """Cancel any in-flight validation that has exceeded its timeout_s.
+
+        Cancelled tasks are moved directly from `_pending_validations` to
+        `_finished_validations` so the MAX_IN_FLIGHT slot is freed immediately
+        and the task flows through the normal finished-processing pipeline
+        without waiting for `ray.wait` to echo the cancellation.
+        """
+        now = time.monotonic()
+        for task, pending in list(self._pending_validations.items()):
+            if (
+                pending.timeout_s is None
+                or now - pending.start_time < pending.timeout_s
+            ):
+                continue
+            self._pending_validations.pop(task)
+            logger.warning(
+                f"Validation for checkpoint {pending.checkpoint} exceeded "
+                f"timeout_s={pending.timeout_s}s. Cancelling."
+            )
+            self._timed_out_tasks.add(task)
+            ray.cancel(task, force=True)
+            self._finished_validations[task] = pending.checkpoint
+
+    def _poll_validations(self) -> int:
+        """Poll/process validations, update checkpoint manager, return num pending validations."""
+        self._cancel_timed_out_validations()
+
+        # Move pending validations to finished validations
+        validation_tasks = list(self._pending_validations.keys())
+        done, _ = ray.wait(
+            validation_tasks, timeout=0, num_returns=len(validation_tasks)
+        )
+        done_checkpoints = []
+        for task in done:
+            pending = self._pending_validations.pop(task)
+            done_checkpoints.append(pending.checkpoint)
+            self._finished_validations[task] = pending.checkpoint
+        if done_checkpoints:
+            logger.info(
+                f"Finished async validation task(s) for checkpoint(s): {done_checkpoints}.\n"
+                f"Running validations for checkpoint(s): {[p.checkpoint for p in self._pending_validations.values()]}.\n"
+                f"Staged validations for checkpoint(s): {[tr.checkpoint for tr in self._training_report_queue]}."
+            )
+
+        # Process finished validations (one at a time)
+        if self._finished_validations:
+            task, checkpoint = self._finished_validations.popitem(last=False)
+            update = self._process_finished_validation(task, checkpoint)
+            if update is not None:
+                self._checkpoint_manager.update_checkpoints_with_validation_result(
+                    {checkpoint: update}
+                )
+
+        return len(self._pending_validations)
+
+    def _kick_off_validations(self) -> int:
+        """Kick off validations and return the number of pending validations."""
+        # TODO: figure out where to place run_validation_fn task:
+        # TODO: provide option to run this on gpu?
+        num_validations_to_start = max(
+            MAX_IN_FLIGHT_VALIDATIONS - len(self._pending_validations), 0
+        )
+        num_validations_to_start = min(
+            num_validations_to_start, len(self._training_report_queue)
+        )
+        for _ in range(num_validations_to_start):
+            training_report = self._training_report_queue.popleft()
+            run_validation_fn_with_options = run_validation_fn.options(
+                **self._validation_config.ray_remote_kwargs,
+            )
+            validate_task = run_validation_fn_with_options.remote(
+                self._validation_config,
+                training_report.validation,
+                training_report.checkpoint,
+            )
+            if isinstance(training_report.validation, ValidationTaskConfig):
+                timeout_s = training_report.validation.timeout_s
+            else:
+                timeout_s = self._validation_config.task_config.timeout_s
+            self._pending_validations[validate_task] = _PendingValidation(
+                checkpoint=training_report.checkpoint,
+                start_time=time.monotonic(),
+                timeout_s=timeout_s,
+            )
+            logger.info(
+                f"Launched async validation task for checkpoint {training_report.checkpoint}"
+            )
+        return len(self._pending_validations)
+
+    def _process_finished_validation(
+        self, task: ray.ObjectRef, checkpoint: Checkpoint
+    ) -> Optional[Tuple[Dict[str, Any], ReportedCheckpointStatus]]:
+        """Process finished validation. Returns (metrics, status) or None.
+
+        Returns None when the task was cancelled by a controller abort (not a
+        timeout), leaving it pending so it re-queues on resumption.
+        """
+        was_timed_out = task in self._timed_out_tasks
+        self._timed_out_tasks.discard(task)
+        if was_timed_out:
+            logger.info(
+                f"Validation for checkpoint {checkpoint} was cancelled due to timeout."
+            )
+            return {}, ReportedCheckpointStatus.VALIDATION_TIMEOUT
+
+        try:
+            metrics = ray.get(task)
+            return metrics, ReportedCheckpointStatus.VALIDATED
+        except ray.exceptions.TaskCancelledError:
+            logger.info(
+                f"Validation was cancelled for checkpoint {checkpoint}, likely because the train run was aborted. "
+                "It will be retried in the next train run with the same storage path if there is one."
+            )
+            return None
+        except ray.exceptions.RayTaskError:
+            logger.exception(f"Validation failed for checkpoint {checkpoint}")
+            return {}, ReportedCheckpointStatus.VALIDATION_FAILED
+
+    async def before_controller_shutdown(self):
+        while self._poll_validations() != 0 or self._kick_off_validations() != 0:
+            await asyncio.sleep(VALIDATION_TASK_POLL_INTERVAL_S)
+        checkpoint_updates: Dict[
+            Checkpoint, Tuple[Dict[str, Any], ReportedCheckpointStatus]
+        ] = {}
+        tasks = list(self._finished_validations.keys())
+        for task in tasks:
+            checkpoint = self._finished_validations[task]
+            self._finished_validations.pop(task)
+            update = self._process_finished_validation(task, checkpoint)
+            if update is not None:
+                checkpoint_updates[checkpoint] = update
+        self._checkpoint_manager.update_checkpoints_with_validation_result(
+            checkpoint_updates
+        )
+
+    def before_controller_abort(self):
+        for task in self._pending_validations.keys():
+            ray.cancel(task)
+
+    def after_controller_state_update(
+        self,
+        previous_state: "TrainControllerState",
+        current_state: "TrainControllerState",
+    ):
+        # TODO: figure out if there's a better place to poll validations
+        if current_state.is_terminal():
+            return
+        self._poll_validations()
+        self._kick_off_validations()
+
+    def before_init_train_context(
+        self, workers: List["Worker"]
+    ) -> Dict[str, List[bool]]:
+        return {
+            "has_validation_fn": [True] * len(workers),
+        }

python/ray/train/v2/_internal/execution/collective_impl.py

@@ -0,0 +1,56 @@
+import logging
+from typing import Any
+
+import ray
+import ray.cloudpickle as pickle
+from ray.train.v2._internal.execution.context import get_train_context
+
+# For reference, {1:1} is 19 bytes, {"1":"1"} is 21 bytes,
+# and {"12345": "12345"} is 25 bytes.
+_MAX_BROADCAST_SIZE_BYTES = 1000
+
+
+logger = logging.getLogger(__name__)
+
+
+def barrier() -> None:
+    """
+    Create a barrier across all training workers.
+    """
+    train_context = get_train_context()
+    sync_actor = train_context.get_synchronization_actor()
+    return ray.get(
+        sync_actor.broadcast_from_rank_zero.remote(
+            world_rank=train_context.get_world_rank(),
+            world_size=train_context.get_world_size(),
+            data=None,
+            caller_method_name="ray.train.collective.barrier",
+        )
+    )
+
+
+def broadcast_from_rank_zero(data: Any) -> Any:
+    """Broadcast data from the rank 0 worker to all other workers.
+
+    This method is used by the public API function :func:`ray.train.collective.broadcast_from_rank_zero`.
+    Users should typically call ``ray.train.collective.broadcast_from_rank_zero()`` instead of calling this method directly.
+    """
+    # Validate data.
+    if data is not None:
+        data_bytes = len(pickle.dumps(data))
+        if data_bytes > _MAX_BROADCAST_SIZE_BYTES:
+            logger.warning(
+                f"Data size {data_bytes} bytes exceeds the maximum broadcast "
+                f"size of {_MAX_BROADCAST_SIZE_BYTES} bytes"
+            )
+
+    train_context = get_train_context()
+    sync_actor = train_context.get_synchronization_actor()
+    return ray.get(
+        sync_actor.broadcast_from_rank_zero.remote(
+            world_rank=train_context.get_world_rank(),
+            world_size=train_context.get_world_size(),
+            data=data,
+            caller_method_name="ray.train.collective.broadcast_from_rank_zero",
+        )
+    )

python/ray/train/v2/_internal/execution/context.py

@@ -0,0 +1,559 @@
+import logging
+import sys
+import threading
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+from pathlib import Path
+from queue import Queue
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+
+import ray
+from ray._common.retry import retry
+from ray._common.utils import env_float
+from ray.actor import ActorHandle
+from ray.train.v2._internal.constants import (
+    AWS_RETRYABLE_TOKENS,
+    CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR,
+    DEFAULT_CHECKPOINT_UPLOAD_WARN_INTERVAL_S,
+)
+from ray.train.v2._internal.execution.checkpoint.sync_actor import (
+    SynchronizationActor,
+    SynchronizationBarrierResetError,
+)
+from ray.train.v2._internal.execution.preemption import PreemptionContext
+from ray.train.v2._internal.execution.storage import StorageContext, delete_fs_path
+from ray.train.v2._internal.execution.training_report import (
+    _TrainingReport,
+)
+from ray.train.v2._internal.util import (
+    construct_user_exception_with_traceback,
+    context_watchdog,
+    invoke_context_managers,
+)
+from ray.train.v2.api.config import RunConfig, ScalingConfig
+from ray.train.v2.api.report_config import (
+    CheckpointConsistencyMode,
+    CheckpointUploadMode,
+)
+from ray.train.v2.api.validation_config import ValidationTaskConfig
+
+if TYPE_CHECKING:
+    from ray.data import DataIterator
+    from ray.train import BackendConfig, Checkpoint, DataConfig
+    from ray.train.v2._internal.data_integration.interfaces import (
+        DatasetShardMetadata,
+        DatasetShardProvider,
+    )
+    from ray.train.v2._internal.execution.callback import TrainContextCallback
+    from ray.train.v2._internal.execution.worker_group.thread_runner import ThreadRunner
+    from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
+
+
+logger = logging.getLogger(__file__)
+
+
+# TODO: make this value manually or automatically configurable.
+MAX_CHECKPOINT_UPLOAD_THREADS = 1
+
+DEFAULT_CHECKPOINT_UPLOAD_WARN_MESSAGE = "Checkpoint upload for {checkpoint_dir_name} has been running for {elapsed}s (warning interval: {interval}s). This may indicate a network issue or slow storage backend. Consider specifying a different filesystem via RunConfig(storage_filesystem=...)."
+CUSTOM_CHECKPOINT_UPLOAD_WARN_MESSAGE = "Custom checkpoint upload for {checkpoint_dir_name} has been running for {elapsed}s (warning interval: {interval}s). This may indicate an issue in your custom upload function passed to `ray.train.report(custom_upload_fn)`."
+
+
+@dataclass(frozen=True)
+class TrainRunContext:
+    """Holds the metadata and context for the current training run."""
+
+    # The unique ID of the training run.
+    run_id: str = field(init=False, default_factory=lambda: uuid.uuid4().hex)
+
+    # The run configuration for the current training run.
+    run_config: RunConfig
+
+    # The configuration passed to the training function.
+    train_loop_config: Optional[Dict]
+
+    # The scaling configuration for the current training run.
+    scaling_config: ScalingConfig
+
+    # The configuration for the training backend (e.g., PyTorch, XGBoost).
+    backend_config: "BackendConfig"
+
+    # The configuration for dataset ingestion and sharding.
+    dataset_config: "DataConfig"
+
+    def get_run_config(self) -> RunConfig:
+        """Returns the run config of the current training run."""
+        return self.run_config
+
+
+@dataclass(frozen=True)
+class DistributedContext:
+    world_rank: int
+    world_size: int
+    local_rank: int
+    local_world_size: int
+    node_rank: int
+
+
+@dataclass(frozen=True)
+class ExecutionContext:
+    """Holds the execution context for the current worker process.
+
+    Every worker process has a single execution context accessed via the
+    `TrainContext`, which includes the training thread that is actually
+    running the user code.
+    """
+
+    # A shared synchronization actor that helps broadcast data across ranks.
+    synchronization_actor: SynchronizationActor
+
+    # A queue that receives training results from the user training code.
+    # `ray.train.report` in user code populates this queue.
+    result_queue: Queue
+
+    # The thread launcher that runs the user training loop.
+    training_thread_runner: "ThreadRunner"
+
+    # The callbacks that are run in the worker train context.
+    train_context_callbacks: List["TrainContextCallback"]
+
+
+@dataclass
+class TrainContext:
+    train_run_context: TrainRunContext
+    distributed_context: DistributedContext
+    execution_context: ExecutionContext
+    storage_context: StorageContext
+    preemption_context: PreemptionContext
+    controller_actor: ActorHandle
+
+    dataset_shard_provider: "DatasetShardProvider"
+    has_validation_fn: Optional[bool] = None
+
+    # TODO: consolidate into CheckpointContext
+    checkpoint: Optional["Checkpoint"] = None
+    current_report_index: int = 0
+    report_call_index: int = 0
+    report_order_condition: threading.Condition = threading.Condition()
+    checkpoint_upload_threadpool: ThreadPoolExecutor = ThreadPoolExecutor(
+        max_workers=MAX_CHECKPOINT_UPLOAD_THREADS
+    )
+
+    def __post_init__(self):
+        # Ray train initializes worker with current report index
+        # report_call_index should start at the current report index
+        self.report_call_index = self.current_report_index
+
+    def get_experiment_name(self) -> str:
+        return self.train_run_context.run_config.name
+
+    def get_world_size(self) -> int:
+        return self.distributed_context.world_size
+
+    def get_world_rank(self) -> int:
+        return self.distributed_context.world_rank
+
+    def get_local_rank(self) -> int:
+        return self.distributed_context.local_rank
+
+    def get_local_world_size(self) -> int:
+        return self.distributed_context.local_world_size
+
+    def get_node_rank(self) -> int:
+        return self.distributed_context.node_rank
+
+    def get_storage(self):
+        return self.storage_context
+
+    # TODO: Don't allow these private methods to be called from user code.
+    def get_result_queue(self):
+        return self.execution_context.result_queue
+
+    def get_synchronization_actor(self):
+        return self.execution_context.synchronization_actor
+
+    def get_checkpoint(self):
+        with self.report_order_condition:
+            return self.checkpoint
+
+    def get_all_reported_checkpoints(
+        self,
+        consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
+        timeout_s: Optional[float] = None,
+    ) -> List["ReportedCheckpoint"]:
+        return ray.get(
+            self.controller_actor.get_all_reported_checkpoints.remote(
+                self.report_call_index,
+                consistency_mode,
+                timeout_s,
+            )
+        )
+
+    def get_dataset_shard(self, dataset_info: "DatasetShardMetadata") -> "DataIterator":
+        """Returns the :class:`ray.data.DataIterator` shard for this worker.
+
+        Call :meth:`~ray.data.DataIterator.iter_torch_batches` or
+        :meth:`~ray.data.DataIterator.to_tf` on this shard to convert it to the
+        appropriate framework-specific data type.
+
+        Args:
+            dataset_info: The shard metadata, including the dataset name and worker rank.
+        Returns:
+            The ``DataIterator`` shard with the given name for this worker.
+        Raises:
+            KeyError: If the dataset shard with the given name is not found.
+        """
+        return self.dataset_shard_provider.get_dataset_shard(dataset_info)
+
+    def get_context_callbacks(self) -> List["TrainContextCallback"]:
+        return self.execution_context.train_context_callbacks
+
+    def _sync_checkpoint_dir_name_across_ranks(
+        self, checkpoint_dir_name: Optional[str] = None
+    ) -> str:
+        """Sync the checkpoint dir name across ranks.
+
+        Args:
+            checkpoint_dir_name: The checkpoint dir name to sync.
+
+        Returns:
+            The synced checkpoint dir name.
+        """
+        # If checkpoint_dir_name is not set, use default checkpoint_dir_name
+        # created by the storage context.
+        checkpoint_dir_name = (
+            checkpoint_dir_name
+            or self.storage_context.make_default_checkpoint_dir_name()
+        )
+        # Get a consensus across ranks on the remote storage path, so distributed
+        # checkpoints will be stored to the same place.
+        sync_actor = self.get_synchronization_actor()
+        with invoke_context_managers(
+            [
+                callback.on_checkpoint_sync
+                for callback in self.execution_context.train_context_callbacks
+            ]
+        ):
+            return ray.get(
+                sync_actor.broadcast_from_rank_zero.remote(
+                    world_rank=self.distributed_context.world_rank,
+                    world_size=self.distributed_context.world_size,
+                    data=checkpoint_dir_name,
+                    caller_method_name="ray.train.report",
+                )
+            )
+
+    # TODO: make retry configurable
+    @retry(description="upload checkpoint", max_attempts=3, match=AWS_RETRYABLE_TOKENS)
+    def _upload_checkpoint(
+        self,
+        checkpoint_dir_name: str,
+        metrics: Dict[str, Any],
+        checkpoint: Optional["Checkpoint"] = None,
+        delete_local_checkpoint_after_upload: bool = False,
+        checkpoint_upload_fn: Optional[
+            Callable[["Checkpoint", str], "Checkpoint"]
+        ] = None,
+        validation: Union[bool, ValidationTaskConfig] = False,
+    ) -> _TrainingReport:
+        """Save the checkpoint to remote storage.
+
+        Args:
+            checkpoint_dir_name: The checkpoint dir to persist to.
+            metrics: The metrics to report.
+            checkpoint: The checkpoint to report.
+            delete_local_checkpoint_after_upload: Whether to delete the checkpoint after it is uploaded.
+            checkpoint_upload_fn: A user defined function that will be called with the
+                checkpoint to upload it. If not provided, defaults to using the `pyarrow.fs.copy_files`
+                utility for copying to the destination `storage_path`.
+            validation: The validation configuration.
+
+        Returns:
+            The training result object containing the persisted checkpoint.
+        """
+
+        if not checkpoint:
+            return _TrainingReport(checkpoint=None, metrics=metrics, validation=False)
+
+        def slow_upload_warning(stop_event: threading.Event, message: str):
+            # Log a warning for the checkpoint upload every `CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR`
+            #   seconds until `stop_event` is set.
+            elapsed = 0.0
+            interval = env_float(
+                CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR,
+                DEFAULT_CHECKPOINT_UPLOAD_WARN_INTERVAL_S,
+            )
+            while not stop_event.wait(interval):
+                elapsed += interval
+                logger.warning(
+                    message.format(
+                        checkpoint_dir_name=checkpoint_dir_name,
+                        elapsed=elapsed,
+                        interval=interval,
+                    )
+                )
+
+        # Records how long the checkpoint transfer took
+        warn_message = (
+            CUSTOM_CHECKPOINT_UPLOAD_WARN_MESSAGE
+            if checkpoint_upload_fn
+            else DEFAULT_CHECKPOINT_UPLOAD_WARN_MESSAGE
+        )
+        with invoke_context_managers(
+            [
+                callback.on_checkpoint_transfer
+                for callback in self.execution_context.train_context_callbacks
+            ]
+        ):
+            try:
+                with context_watchdog(slow_upload_warning, warn_message):
+                    if checkpoint_upload_fn:
+                        # Upload the checkpoint using the custom checkpoint_upload_fn
+                        persisted_checkpoint = checkpoint_upload_fn(
+                            checkpoint, checkpoint_dir_name
+                        )
+                    else:
+                        # Upload the checkpoint using PyArrow
+                        persisted_checkpoint = (
+                            self.storage_context.persist_current_checkpoint(
+                                checkpoint, checkpoint_dir_name
+                            )
+                        )
+            except FileNotFoundError:
+                logger.exception(
+                    f"Failed to find local checkpoint ({checkpoint}) when attempting to upload it. "
+                    "This could be caused by multiple workers on a node attempting to upload the "
+                    "same directory, and then one of the workers deletes the directory before the "
+                    "others finish."
+                )
+                raise
+
+        # Check that the checkpoint generated is a `ray.train.Checkpoint` instance
+        if checkpoint_upload_fn and not isinstance(
+            persisted_checkpoint, ray.train.Checkpoint
+        ):
+            raise ValueError(
+                f"checkpoint_upload_fn must return a `ray.train.Checkpoint`. Actual type is {type(persisted_checkpoint)}"
+            )
+
+        # TODO: consider deleting local checkpoint as async callback instead
+        if delete_local_checkpoint_after_upload:
+            try:
+                delete_fs_path(checkpoint.filesystem, checkpoint.path)
+            except Exception:
+                logger.exception(
+                    f"Failed to delete the local checkpoint after a successful upload: {checkpoint}"
+                )
+
+        return _TrainingReport(
+            checkpoint=persisted_checkpoint,
+            metrics=metrics,
+            validation=validation,
+        )
+
+    def _wait_then_report(
+        self, training_report: _TrainingReport, report_call_index: int
+    ):
+        """Thread waits for its turn before reporting training result to result queue.
+
+        It does this in order to guarantee the FIFO processing of checkpoints.
+
+        The queue size is set to 1 to avoid accumulating unprocessed results.
+        If the queue is full, the put operation blocks until a result is consumed.
+
+        TODO: Add a metric to track the blocking time waiting for the
+        training result to be consumed by the controller.
+        """
+        with self.report_order_condition:
+            self.report_order_condition.wait_for(
+                lambda: self.current_report_index == report_call_index - 1
+            )
+            logger.info(
+                f"Reporting training result {report_call_index}: {training_report} "
+                f"from rank {self.get_world_rank()}"
+            )
+            # Update latest checkpoint as the persisted checkpoint.
+            if training_report.checkpoint:
+                self.checkpoint = training_report.checkpoint
+            self.get_result_queue().put(training_report)
+            self.current_report_index += 1
+            self.report_order_condition.notify_all()
+
+    def report(
+        self,
+        metrics: Dict[str, Any],
+        checkpoint: Optional["Checkpoint"] = None,
+        checkpoint_dir_name: Optional[str] = None,
+        checkpoint_upload_mode: CheckpointUploadMode = CheckpointUploadMode.SYNC,
+        delete_local_checkpoint_after_upload: Optional[bool] = None,
+        checkpoint_upload_fn: Optional[
+            Callable[["Checkpoint", str], "Checkpoint"]
+        ] = None,
+        validation: Union[bool, ValidationTaskConfig] = False,
+    ) -> None:
+        """
+        Upload checkpoint to remote storage and put a training
+        result on the result queue of this worker process.
+
+        TODO: the report function should be implemented in the worker instead
+          of in the train context. The train context should only keep the train
+          related information and not the worker related actions. This refactor
+          would also require the `TrainContextCallback` to be updated as well.
+        """
+        if "torch" in sys.modules:
+            from ray.air._internal.torch_utils import contains_tensor
+
+            if contains_tensor(metrics):
+                raise ValueError(
+                    "Passing objects containing Torch tensors as metrics "
+                    "is not supported as it will throw an exception on "
+                    "deserialization. You can either convert the tensors "
+                    "to Python objects (ex: `.numpy()`, `.item()`, etc.) "
+                    "or save tensors as part of the checkpoint files instead."
+                )
+
+        if validation and not self.has_validation_fn:
+            raise ValueError(
+                "`validation_config` was not set on the trainer, but a validation was requested."
+            )
+
+        if delete_local_checkpoint_after_upload and checkpoint is not None:
+            experiment_path = Path(self.storage_context.experiment_fs_path)
+            checkpoint_path = Path(checkpoint.path)
+
+            # Resolve symlinks only for local (absolute) paths.
+            # Remote paths (S3, GCS, etc.) are relative after URI and resolve()
+            # would prepend CWD, producing a meaningless local path.
+            # Mixed absolute/relative paths return False
+            if experiment_path.is_absolute():
+                experiment_path = experiment_path.resolve()
+            if checkpoint_path.is_absolute():
+                checkpoint_path = checkpoint_path.resolve()
+
+            if experiment_path.is_relative_to(checkpoint_path):
+                raise ValueError(
+                    f"Ray Train's experiment directory ({self.storage_context.experiment_fs_path}) "
+                    f"is contained within the checkpoint path ({checkpoint.path}) "
+                    f"and `ray.train.report(delete_local_checkpoint_after_upload=True)`. "
+                    "As a result, this would delete the experiment directory. "
+                    "Please write the checkpoint to a temporary directory, "
+                    "a subdirectory of the experiment directory, "
+                    "or use `delete_local_checkpoint_after_upload=False`."
+                )
+
+        with invoke_context_managers(
+            [
+                callback.on_report
+                for callback in self.execution_context.train_context_callbacks
+            ]
+        ):
+            self.report_call_index += 1
+            report_call_index = self.report_call_index
+
+            # Sync the checkpoint dir name across ranks.
+            try:
+                checkpoint_dir_name = self._sync_checkpoint_dir_name_across_ranks(
+                    checkpoint_dir_name
+                )
+            except ray.exceptions.RayTaskError as e:
+                if not isinstance(e.cause, SynchronizationBarrierResetError):
+                    raise e
+                logger.warning(
+                    "Synchronization barrier was reset (likely due to a "
+                    "worker failure). Skipping this report."
+                )
+                # Keep report indexes aligned across workers.
+                self.report_call_index -= 1
+                return
+
+            # Upload checkpoint, wait for turn, and report.
+            if checkpoint_upload_mode == CheckpointUploadMode.SYNC:
+                training_report = self._upload_checkpoint(
+                    checkpoint_dir_name,
+                    metrics,
+                    checkpoint,
+                    delete_local_checkpoint_after_upload,
+                    checkpoint_upload_fn,
+                    validation,
+                )
+                self._wait_then_report(training_report, report_call_index)
+
+            elif checkpoint_upload_mode == CheckpointUploadMode.NO_UPLOAD:
+                training_report = _TrainingReport(
+                    checkpoint=checkpoint,
+                    metrics=metrics,
+                    validation=validation,
+                )
+                self._wait_then_report(training_report, report_call_index)
+
+            elif checkpoint_upload_mode == CheckpointUploadMode.ASYNC:
+
+                def _upload_checkpoint_and_report(
+                    checkpoint_dir_name: str,
+                    metrics: Dict[str, Any],
+                    checkpoint: Optional["Checkpoint"],
+                    report_call_index: int,
+                ) -> None:
+                    try:
+                        training_report = self._upload_checkpoint(
+                            checkpoint_dir_name,
+                            metrics,
+                            checkpoint,
+                            delete_local_checkpoint_after_upload,
+                            checkpoint_upload_fn,
+                            validation,
+                        )
+                        self._wait_then_report(training_report, report_call_index)
+                    except Exception as e:
+                        # TODO: env var to disable eager raising
+                        logger.exception(
+                            "Checkpoint upload failed in the background thread. Raising eagerly "
+                            "to avoid training in a corrupted state with more potential progress "
+                            "lost due to checkpointing failures."
+                        )
+                        self.execution_context.training_thread_runner.get_exception_queue().put(
+                            construct_user_exception_with_traceback(e)
+                        )
+
+                self.checkpoint_upload_threadpool.submit(
+                    _upload_checkpoint_and_report,
+                    checkpoint_dir_name,
+                    metrics,
+                    checkpoint,
+                    report_call_index,
+                )
+            else:
+                raise ValueError(
+                    f"Invalid checkpoint upload mode: {checkpoint_upload_mode}"
+                )
+
+
+# The global variable holding the current TrainContext
+_train_context: Optional[TrainContext] = None
+
+# Thread lock to protect the global TrainContext
+_context_lock = threading.Lock()
+
+
+def get_train_context() -> TrainContext:
+    """Get the internal train context.
+
+    Note:
+        This should not be used directly by user-facing APIs. User-facing APIs should
+        call :class:`~ray.train.v2._internal.execution.train_fn_utils.TrainFnUtils`
+        or use :class:`~ray.train.v2.api.context.TrainContext` instead.
+
+    Returns:
+        The internal TrainContext for this worker.
+    """
+    with _context_lock:
+        if _train_context is None:
+            raise RuntimeError("TrainContext has not been initialized.")
+        return _train_context
+
+
+def set_train_context(context) -> None:
+    global _train_context
+    with _context_lock:
+        _train_context = context

python/ray/train/v2/_internal/execution/controller/__init__.py

@@ -0,0 +1,3 @@
+from .controller import TrainController
+
+__all__ = ["TrainController"]

python/ray/train/v2/_internal/execution/controller/controller.py

@@ -0,0 +1,865 @@
+import asyncio
+import logging
+import os
+import uuid
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, List, Optional
+
+import pandas as pd
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray.exceptions import AsyncioActorExit
+from ray.train.v2._internal.constants import (
+    DEFAULT_ENABLE_CONTROLLER_LOGGING,
+    DEFAULT_ENABLE_PREEMPTION_WATCHER,
+    DEFAULT_HEALTH_CHECK_INTERVAL_S,
+    ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR,
+    ENABLE_PREEMPTION_WATCHER_ENV_VAR,
+    HEALTH_CHECK_INTERVAL_S_ENV_VAR,
+)
+from ray.train.v2._internal.execution.callback import (
+    ControllerCallback,
+    ReportCallback,
+    TrainContextCallback,
+    WorkerCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.callback_manager import CallbackManager
+from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
+    CheckpointManager,
+)
+from ray.train.v2._internal.execution.checkpoint.report_handler import (
+    ReportCallbackHandler,
+)
+from ray.train.v2._internal.execution.checkpoint.validation_manager import (
+    ValidationManager,
+)
+from ray.train.v2._internal.execution.context import TrainRunContext
+from ray.train.v2._internal.execution.controller.state import (
+    AbortedState,
+    ErroredState,
+    FinishedState,
+    InitializingState,
+    ReschedulingState,
+    ResizingState,
+    RestartingState,
+    RunningState,
+    SchedulingState,
+    ShuttingDownState,
+    TrainControllerState,
+)
+from ray.train.v2._internal.execution.failure_handling import (
+    FailureDecision,
+    FailurePolicy,
+)
+from ray.train.v2._internal.execution.scaling_policy import (
+    NoopDecision,
+    ResizeDecision,
+    ScalingPolicy,
+)
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroup,
+    WorkerGroupContext,
+    WorkerGroupPollStatus,
+)
+from ray.train.v2._internal.logging import LoggingManager
+from ray.train.v2._internal.util import ObjectRefWrapper, time_monotonic
+from ray.train.v2.api.callback import RayTrainCallback
+from ray.train.v2.api.exceptions import (
+    ControllerError,
+    TrainingFailedError,
+)
+from ray.train.v2.api.report_config import CheckpointConsistencyMode
+from ray.train.v2.api.result import Result
+from ray.train.v2.api.validation_config import ValidationConfig
+
+if TYPE_CHECKING:
+    from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
+
+from ray.util.tpu import get_tpu_num_slices_for_workers
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TrainControllerLoopIterationResult:
+    """The result of a single iteration of the control loop."""
+
+    run_attempt_id: str
+    previous_state: TrainControllerState
+    next_state: TrainControllerState
+    training_failed_error: Optional[TrainingFailedError] = None
+
+    def __repr__(self) -> str:
+        return (
+            f"TrainControllerLoopIterationResult(\n"
+            f"    run_attempt_id={self.run_attempt_id},\n"
+            f"    previous_state={self.previous_state._state_type.state_name},\n"
+            f"    next_state={self.next_state._state_type.state_name}\n"
+            f"    training_failed_error={self.training_failed_error}\n"
+            f")"
+        )
+
+
+class TrainController:
+    """Manages the execution of a distributed training job.
+
+    Responsibilities include:
+    * Triggering the training function to run on the worker group.
+    * Monitoring the status of the worker group.
+    * Handling scaling decisions by restarting the worker group.
+    * Handling failure decisions by restarting the worker group or terminating training.
+    * Running callback logic on different hooks in the control loop.
+    """
+
+    worker_group_cls = WorkerGroup
+
+    def __init__(
+        self,
+        train_fn_ref: ObjectRefWrapper[Callable[[], None]],
+        train_run_context: TrainRunContext,
+        scaling_policy: ScalingPolicy,
+        failure_policy: FailurePolicy,
+        callbacks: Optional[List[RayTrainCallback]] = None,
+        validation_config: Optional[ValidationConfig] = None,
+    ):
+        self._train_run_context = train_run_context
+        if ray_constants.env_bool(
+            ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR,
+            DEFAULT_ENABLE_CONTROLLER_LOGGING,
+        ):
+            LoggingManager.configure_controller_logger(self._train_run_context)
+        self._train_fn_ref = train_fn_ref
+        self._scaling_policy = scaling_policy
+        self._failure_policy = failure_policy
+        self._run_config = self._train_run_context.run_config
+        self._callbacks = callbacks or []
+        self._storage_context = self._train_run_context.run_config.storage_context
+
+        self._checkpoint_manager = CheckpointManager(
+            checkpoint_config=self._run_config.checkpoint_config,
+            storage_context=self._storage_context,
+        )
+        if validation_config:
+            validation_manager = ValidationManager(
+                checkpoint_manager=self._checkpoint_manager,
+                validation_config=validation_config,
+            )
+        else:
+            validation_manager = None
+        report_handler = ReportCallbackHandler(
+            report_callbacks=(
+                [self._checkpoint_manager]
+                + ([validation_manager] if validation_manager else [])
+                + [c for c in self._callbacks if isinstance(c, ReportCallback)]
+            )
+        )
+
+        # Group callbacks by the hooks they're subscribed to.
+        self._controller_callbacks = (
+            [
+                self._scaling_policy,
+            ]
+            + ([validation_manager] if validation_manager else [])
+            + [c for c in self._callbacks if isinstance(c, ControllerCallback)]
+        )
+        self._controller_callback_manager = CallbackManager(self._controller_callbacks)
+        # Group callbacks that will be propagated to the worker group,
+        # train worker and the train context.
+        self._worker_group_callbacks_to_propagate = (
+            [report_handler]
+            + ([validation_manager] if validation_manager else [])
+            + [
+                c
+                for c in self._callbacks
+                if isinstance(
+                    c, (WorkerGroupCallback, WorkerCallback, TrainContextCallback)
+                )
+            ]
+            + [self._checkpoint_manager]
+        )
+
+        self._health_check_interval_s = float(
+            os.getenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, DEFAULT_HEALTH_CHECK_INTERVAL_S)
+        )
+
+        self._manages_replica_groups = (
+            train_run_context.backend_config.backend_cls.has_replica_groups
+            if train_run_context.backend_config
+            else False
+        )
+
+        # Register the preemption-observability callback when not in TorchFT
+        # mode (replica groups handle peer loss via their own quorum).
+        enable_preemption_watcher = ray_constants.env_bool(
+            ENABLE_PREEMPTION_WATCHER_ENV_VAR,
+            DEFAULT_ENABLE_PREEMPTION_WATCHER,
+        )
+        if self._manages_replica_groups:
+            if enable_preemption_watcher and ray_constants.env_set_by_user(
+                ENABLE_PREEMPTION_WATCHER_ENV_VAR
+            ):
+                logger.info(
+                    "The preemption watcher is not compatible with replica "
+                    "groups (e.g. TorchFT), which handle peer loss via their "
+                    "own quorum; skipping it."
+                )
+        elif enable_preemption_watcher:
+            from ray.train.v2._internal.callbacks.preemption_callback import (
+                PreemptionCallback,
+            )
+
+            self._worker_group_callbacks_to_propagate.append(PreemptionCallback())
+
+        self._worker_group: Optional[WorkerGroup] = None
+        self._state = InitializingState()
+        self._return_value: Optional[Any] = None
+
+        # TODO: These can be attributes of a RunAttempt?
+        self._latest_poll_time = float("-inf")
+
+        # Generate an initial run attempt ID so that `_run_controller_hook`
+        # can reference it if a callback fails during `_start`.
+        self._generate_run_attempt_id()
+        self._start()
+
+    def _run_controller_hook(
+        self,
+        hook_name: str,
+        *args,
+        invoke_failure_decision_callbacks: bool = True,
+        **context,
+    ) -> Optional["TrainControllerLoopIterationResult"]:
+        """Invoke a named controller hook and catch any exceptions.
+
+        This method invokes all callbacks registered for the given controller hook.
+        If a callback raises an error, the error is routed through the failure policy
+        and may produce a ``TrainControllerLoopIterationResult``, indicating that the
+        current controller step should exit early with this failure result.
+
+        Args:
+            hook_name: The controller hook name to invoke.
+            *args: Positional arguments to pass to the hook.
+            invoke_failure_decision_callbacks: Whether to invoke failure-decision hooks
+                when handling a callback failure.
+            **context: Keyword arguments to pass to the hook.
+
+        Returns:
+            failure_result: A``TrainControllerLoopIterationResult`` if the hook execution results
+            in an early exit from the controller loop to raise the callback error,
+            or ``None`` if hook execution completes successfully.
+        """
+        try:
+            self._controller_callback_manager.invoke(hook_name, *args, **context)
+        except ControllerError as error:
+            failure_decision = self._failure_policy.make_decision(
+                training_failed_error=error,
+            )
+            # Avoid re-entering controller callback hooks while handling a callback failure.
+            return self._execute_failure_decision(
+                failure_decision,
+                training_failed_error=error,
+                invoke_failure_decision_callbacks=invoke_failure_decision_callbacks,
+            )
+        return None
+
+    def _execute_resize_decision(
+        self, decision: ResizeDecision
+    ) -> TrainControllerLoopIterationResult:
+        """Executes resize decisions.
+
+        Errors from worker group shutdown, callbacks, or worker group startup
+        are allowed to propagate to the catch-all in ``run()``.
+        """
+
+        failure_result = self._run_controller_hook(
+            "before_controller_execute_resize_decision", decision
+        )
+        if failure_result:
+            return failure_result
+        current_num_workers = (
+            len(self._worker_group.get_workers()) if self._worker_group else 0
+        )
+        poll_status = (
+            self._worker_group.get_latest_poll_status() if self._worker_group else None
+        )
+        failing_rgs = (
+            poll_status.failing_replica_group_indices if poll_status else set()
+        )
+        all_rgs = poll_status.all_replica_group_indices if poll_status else set()
+        if (
+            self._manages_replica_groups
+            and bool(failing_rgs)
+            and failing_rgs != all_rgs
+            and self._worker_group
+            # TODO: relax this after integrating replica groups with elastic training.
+            and decision.num_workers == current_num_workers
+        ):
+            # Torchft: replace only failing replica groups.
+            self._replace_bad_workers(poll_status)
+        else:
+            # Standard: full restart.
+            if self._worker_group:
+                self._shutdown_worker_group()
+            self._start_worker_group(
+                num_workers=decision.num_workers,
+                resources_per_worker=decision.resources_per_worker,
+            )
+
+        return TrainControllerLoopIterationResult(
+            run_attempt_id=self._get_run_attempt_id(),
+            previous_state=self._state,
+            next_state=RunningState(),
+        )
+
+    def _replace_bad_workers(self, poll_status: WorkerGroupPollStatus):
+        """Replace failing replica groups in the worker group.
+
+        Args:
+            poll_status: The poll status containing error information.
+
+        Returns:
+            None
+        """
+        failing_rg_indices = poll_status.failing_replica_group_indices
+
+        if not failing_rg_indices:
+            logger.warning("No failing replica groups found in poll status.")
+            return
+
+        logger.info(f"Replacing failing replica groups: {failing_rg_indices}")
+
+        for rg_index in failing_rg_indices:
+            # TODO: parallelize this.
+            # TODO: also ensure that if earlier replacements succeed and later replacements fail,
+            # we don't redo the earlier replacements.
+            # See https://github.com/ray-project/ray/pull/61475#discussion_r3055217289
+            self._worker_group.replace_replica_group(rg_index)
+
+    def _get_retry_state(
+        self,
+        controller_state: TrainControllerState,
+        training_failed_error: TrainingFailedError,
+    ) -> TrainControllerState:
+        if isinstance(controller_state, RunningState):
+            return RestartingState(training_failed_error=training_failed_error)
+        elif isinstance(controller_state, SchedulingState):
+            return ReschedulingState(training_failed_error=training_failed_error)
+        else:
+            # Cannot retry from this state (e.g. InitializingState,
+            # ShuttingDownState); force shutdown with error.
+            logger.warning(
+                "Cannot retry from state %s; forcing shutdown.",
+                type(controller_state).__name__,
+            )
+            return ShuttingDownState(
+                next_state=ErroredState(training_failed_error=training_failed_error)
+            )
+
+    def _execute_failure_decision(
+        self,
+        failure_decision: FailureDecision,
+        training_failed_error: TrainingFailedError,
+        invoke_failure_decision_callbacks: bool = True,
+    ) -> TrainControllerLoopIterationResult:
+        """Executes failure handling decisions for a scheduling or poll error."""
+
+        controller_state = self.get_state()
+
+        if invoke_failure_decision_callbacks:
+            failure_result = self._run_controller_hook(
+                "before_controller_execute_failure_decision",
+                failure_decision,
+                invoke_failure_decision_callbacks=False,
+            )
+            if failure_result:
+                return failure_result
+
+        # TODO: What should we do here?
+        # This currently never happens because there must be errors.
+        if failure_decision == FailureDecision.NOOP:
+            return TrainControllerLoopIterationResult(
+                run_attempt_id=self._get_run_attempt_id(),
+                previous_state=controller_state,
+                next_state=controller_state,
+                training_failed_error=training_failed_error,
+            )
+
+        if failure_decision == FailureDecision.RETRY:
+            return TrainControllerLoopIterationResult(
+                run_attempt_id=self._get_run_attempt_id(),
+                previous_state=controller_state,
+                next_state=self._get_retry_state(
+                    controller_state, training_failed_error
+                ),
+            )
+        elif failure_decision == FailureDecision.RAISE:
+            next_state = ShuttingDownState(
+                next_state=ErroredState(
+                    training_failed_error=training_failed_error,
+                ),
+            )
+            return TrainControllerLoopIterationResult(
+                run_attempt_id=self._get_run_attempt_id(),
+                previous_state=controller_state,
+                next_state=next_state,
+                training_failed_error=training_failed_error,
+            )
+        else:
+            raise ValueError(f"Unexpected failure decision: {failure_decision}")
+
+    async def _poll_workers(self) -> WorkerGroupPollStatus:
+        # Ensure that the time between polls is at least HEALTH_CHECK_INTERVAL_S.
+        time_since_last_poll = time_monotonic() - self._latest_poll_time
+        if time_since_last_poll < self._health_check_interval_s:
+            remaining_time = max(
+                self._health_check_interval_s - time_since_last_poll, 0
+            )
+            await asyncio.sleep(remaining_time)
+            if self.get_state().is_terminal():
+                logger.debug(
+                    f"Controller is unexpectedly in terminal state {self.get_state()} after "
+                    "sleeping and before polling workers. Exiting actor."
+                )
+                ray.actor.exit_actor()
+
+        status = self._worker_group.poll_status(timeout=self._health_check_interval_s)
+        self._latest_poll_time = time_monotonic()
+        return status
+
+    def _start_worker_group(self, num_workers: int, resources_per_worker: dict) -> None:
+        """Start the worker group and launch the train function.
+
+        Args:
+            num_workers: The number of workers to start.
+            resources_per_worker: The resources per worker to start.
+
+        Raises:
+            Exception: If the worker group failed to start.
+        """
+        placement_strategy = self._scaling_policy.scaling_config.placement_strategy
+        scaling_config = self._train_run_context.scaling_config
+
+        # Check for `label_selector` to influence WorkerGroup scheduling.
+        label_selector = scaling_config._label_selector_per_worker(num_workers)
+        for callback in self._controller_callbacks:
+            selector = callback.on_controller_start_worker_group(
+                scaling_config=scaling_config, num_workers=num_workers
+            )
+            if selector:
+                if label_selector:
+                    logger.warning(
+                        f"Overriding `ScalingConfig.label_selector` {label_selector} "
+                        f"with label_selector returned by user-specified callback {selector}"
+                    )
+                label_selector = [selector.copy() for _ in range(num_workers)]
+
+        # Calculate num_slices for the worker group if using TPU.
+        num_slices = 1
+        if scaling_config.use_tpu:
+            num_slices = get_tpu_num_slices_for_workers(
+                topology=scaling_config.topology,
+                accelerator_type=scaling_config.accelerator_type,
+                num_workers=num_workers,
+                resources_per_worker=resources_per_worker,
+            )
+
+        worker_group_context = WorkerGroupContext(
+            run_attempt_id=self._get_run_attempt_id(),
+            train_fn_ref=self._train_fn_ref,
+            num_workers=num_workers,
+            resources_per_worker=resources_per_worker,
+            placement_strategy=placement_strategy,
+            label_selector=label_selector,
+            num_slices=num_slices,
+        )
+        self._worker_group = self.worker_group_cls.create(
+            train_run_context=self._train_run_context,
+            worker_group_context=worker_group_context,
+            callbacks=self._worker_group_callbacks_to_propagate,
+        )
+
+    def _start(self):
+        failure_result = self._run_controller_hook(
+            "after_controller_start", self._train_run_context
+        )
+        if failure_result:
+            self._set_state(failure_result.next_state)
+
+    async def _shutdown(self) -> "TrainControllerLoopIterationResult":
+        """Execute shutdown and return the final state transition.
+
+        Shutdown errors are never retried. If an error occurs during shutdown:
+        - If we're already shutting down after a training error
+          (next_state is ErroredState), the original error is preserved.
+        - Otherwise the shutdown error becomes the training failure.
+        """
+        controller_state = self.get_state()
+        assert isinstance(controller_state, ShuttingDownState)
+
+        shutdown_error = None
+
+        # TODO: move to __del__ after https://github.com/ray-project/ray/issues/53169
+        if self._worker_group:
+            try:
+                self._shutdown_worker_group()
+            except Exception as e:
+                logger.exception("Error shutting down worker group.")
+                shutdown_error = ControllerError(e)
+
+        try:
+            await self._controller_callback_manager.async_invoke(
+                "before_controller_shutdown"
+            )
+        except ControllerError as e:
+            if shutdown_error:
+                logger.warning(
+                    "An additional error occurred in the before_controller_shutdown "
+                    "callback after a worker group shutdown error. "
+                    "This error is being ignored to preserve the original "
+                    "shutdown error. Error: %s",
+                    e,
+                )
+            else:
+                shutdown_error = e
+
+        if shutdown_error:
+            if isinstance(controller_state.next_state, ErroredState):
+                logger.warning(
+                    "Another error occurred during shutdown after a training error. "
+                    "This error is being ignored to preserve the original "
+                    "training error. Error: %s",
+                    shutdown_error,
+                )
+            else:
+                return TrainControllerLoopIterationResult(
+                    run_attempt_id=self._get_run_attempt_id(),
+                    previous_state=controller_state,
+                    next_state=ErroredState(training_failed_error=shutdown_error),
+                    training_failed_error=shutdown_error,
+                )
+
+        return TrainControllerLoopIterationResult(
+            run_attempt_id=self._get_run_attempt_id(),
+            previous_state=controller_state,
+            next_state=controller_state.next_state,
+        )
+
+    def _shutdown_worker_group(self):
+        """Shutdown the worker group and set the worker group to None."""
+        self._worker_group.shutdown()
+        self._worker_group = None
+
+    def get_worker_group(self) -> Optional[WorkerGroup]:
+        return self._worker_group
+
+    def get_state(self) -> TrainControllerState:
+        return self._state
+
+    def _set_state(self, state: TrainControllerState):
+        previous_state = self._state
+        self._state = state
+
+        failure_result = self._run_controller_hook(
+            "after_controller_state_update", previous_state, state
+        )
+        if failure_result:
+            # If we're transitioning into a terminal state, or if we're already in the shutdown path to an errored terminal state
+            # (ShuttingDownState -> ErroredState), preserve the original failure as the
+            # surfaced error. A failure in a state-update callback should not overwrite
+            # the underlying root-cause error.
+            if state.is_terminal() or (
+                isinstance(state, ShuttingDownState)
+                and isinstance(state.next_state, ErroredState)
+            ):
+                logger.warning(
+                    "A callback failed during a terminal state transition. "
+                    "This failure is being ignored to preserve the original "
+                    "training result. Error: %s",
+                    failure_result.training_failed_error,
+                )
+                return
+
+            # NOTE: We intentionally do *not* re-invoke `after_controller_state_update`
+            # for this transition to avoid re-entering callback hooks while handling
+            # a callback failure.
+            self._state = failure_result.next_state
+
+    def _make_and_handle_scaling_decision_for_non_running_worker_group(
+        self,
+        controller_state: TrainControllerState,
+    ) -> TrainControllerLoopIterationResult:
+        """Make a scaling decision for a non-running worker group and return the appropriate next state.
+
+        This method should be called when entering a state that requires a scaling decision
+        for a non-running worker group.
+
+        This method handles the complete flow of:
+        1. Shutting down the non-running worker group if it still exists.
+        2. Getting a scaling decision for a non-running worker group
+        3. Determining the next state based on the decision type
+        4. Creating and returning the iteration result
+
+        Args:
+            controller_state: The current controller state
+
+        Returns:
+            TrainControllerLoopIterationResult with the appropriate next state
+        """
+        scaling_decision = (
+            self._scaling_policy.make_decision_for_non_running_worker_group()
+        )
+
+        if isinstance(scaling_decision, NoopDecision):
+            next_state = controller_state
+        elif isinstance(scaling_decision, ResizeDecision):
+            next_state = SchedulingState(scaling_decision)
+        else:
+            raise ValueError(f"Unexpected scaling decision: {scaling_decision}")
+
+        return TrainControllerLoopIterationResult(
+            run_attempt_id=self._get_run_attempt_id(),
+            previous_state=controller_state,
+            next_state=next_state,
+        )
+
+    async def _step(self) -> TrainControllerLoopIterationResult:
+        """Run a single iteration of the control loop.
+
+        Returns:
+            The result of the iteration.
+        """
+        controller_state = self.get_state()
+
+        if isinstance(
+            controller_state, (InitializingState, RestartingState, ReschedulingState)
+        ):
+            return self._make_and_handle_scaling_decision_for_non_running_worker_group(
+                controller_state
+            )
+        elif isinstance(controller_state, SchedulingState):
+            assert isinstance(controller_state.scaling_decision, ResizeDecision)
+            return self._execute_resize_decision(controller_state.scaling_decision)
+        elif isinstance(controller_state, RunningState):
+            worker_group_status: WorkerGroupPollStatus = await self._poll_workers()
+
+            if worker_group_status.finished and not worker_group_status.errors:
+                self._return_value = worker_group_status.worker_statuses[0].return_value
+                return TrainControllerLoopIterationResult(
+                    run_attempt_id=self._get_run_attempt_id(),
+                    previous_state=controller_state,
+                    next_state=ShuttingDownState(
+                        next_state=FinishedState(),
+                    ),
+                )
+            if worker_group_status.errors:
+                worker_group_error = worker_group_status.get_worker_group_error()
+                failure_decision = self._failure_policy.make_decision(
+                    training_failed_error=worker_group_error,
+                )
+                return self._execute_failure_decision(
+                    failure_decision, training_failed_error=worker_group_error
+                )
+
+            scaling_decision = (
+                self._scaling_policy.make_decision_for_running_worker_group(
+                    worker_group_state=self.get_worker_group().get_worker_group_state(),
+                    worker_group_status=worker_group_status,
+                )
+            )
+
+            if isinstance(scaling_decision, NoopDecision):
+                next_state = RunningState()
+            elif isinstance(scaling_decision, ResizeDecision):
+                next_state = ResizingState(
+                    scaling_decision=scaling_decision,
+                )
+            else:
+                raise ValueError(f"Unexpected scaling decision: {scaling_decision}")
+
+            return TrainControllerLoopIterationResult(
+                run_attempt_id=self._get_run_attempt_id(),
+                previous_state=controller_state,
+                next_state=next_state,
+            )
+        elif isinstance(controller_state, ResizingState):
+            return TrainControllerLoopIterationResult(
+                run_attempt_id=self._get_run_attempt_id(),
+                previous_state=controller_state,
+                next_state=SchedulingState(
+                    scaling_decision=controller_state.scaling_decision
+                ),
+            )
+        elif isinstance(controller_state, ShuttingDownState):
+            return await self._shutdown()
+        else:
+            raise ValueError(f"Unexpected controller state: {controller_state}")
+
+    def _generate_run_attempt_id(self):
+        self._run_attempt_id = uuid.uuid4().hex
+        return self._run_attempt_id
+
+    def _get_run_attempt_id(self):
+        return self._run_attempt_id
+
+    async def _run_control_loop_iteration(self):
+        """Run a single iteration of the control loop.
+
+        Steps:
+        1. Poll the worker group for status.
+        2. If the worker group is initializing or recovering from an error,
+            make a scaling decision and execute it.
+        3. If the worker group has finished, set the controller state to FINISHED.
+        4. If the worker group has errors, make a failure decision and execute it.
+        5. Otherwise, the worker group is running healthily.
+            Query the scaling policy for a scaling decision and execute it.
+
+        Errors raised by ``_step`` are caught and routed through the failure
+        policy (retry / raise).  If the failure policy itself fails, the
+        controller is forced into ``ErroredState`` as a last resort.
+
+        ``AsyncioActorExit`` is always re-raised so that the actor can shut
+        down cleanly.
+        """
+        controller_state = self.get_state()
+        assert not controller_state.is_terminal()
+
+        if controller_state.needs_new_run_attempt():
+            self._generate_run_attempt_id()
+
+        try:
+            result = await self._step()
+        except AsyncioActorExit:
+            raise
+        except Exception as e:
+            # Preserve the original error type if it is already a
+            # TrainingFailedError (e.g. WorkerGroupError); otherwise
+            # wrap it in a ControllerError.
+            if isinstance(e, TrainingFailedError):
+                training_error = e
+            else:
+                # Log the full traceback only for unexpected errors.
+                logger.exception("Error in control loop iteration: %s", e)
+                training_error = ControllerError(e)
+            try:
+                failure_decision = self._failure_policy.make_decision(
+                    training_failed_error=training_error,
+                )
+                result = self._execute_failure_decision(
+                    failure_decision,
+                    training_failed_error=training_error,
+                )
+            except Exception:
+                # Last resort: force into errored state, bypassing callbacks.
+                logger.exception(
+                    "Failed to execute failure decision, forcing error state."
+                )
+                self._state = ErroredState(training_failed_error=training_error)
+                return
+
+        self._set_state(result.next_state)
+
+    async def run(self):
+        """Run the main control loop. Exits when training is finished or errored."""
+        while not self.get_state().is_terminal():
+            await self._run_control_loop_iteration()
+
+        # Call after_controller_finish with the final result.
+        result = self._build_result()
+        failure_result = self._run_controller_hook(
+            "after_controller_finish", result, invoke_failure_decision_callbacks=False
+        )
+        # Since we are already in a terminal state, a callback failure should
+        # not overwrite the training outcome — log and preserve the result.
+        if failure_result:
+            logger.warning(
+                "A callback failed after training finished. "
+                "This failure is being ignored to preserve the original "
+                "training result. Error: %s",
+                failure_result.training_failed_error,
+            )
+
+    async def abort(self):
+        """Trigger callback abort hooks and terminate the controller process."""
+        # Do not abort run if it's already finished.
+        if self.get_state().is_terminal():
+            return
+
+        self._controller_callback_manager.invoke_best_effort("before_controller_abort")
+
+        # Intentionally abort worker group before setting train run state because
+        # we only reconcile the states of live train runs.
+        try:
+            if self._worker_group:
+                self._worker_group.abort()
+            self._set_state(AbortedState())
+        except Exception as e:
+            logger.exception("Error aborting worker group: %s", e)
+
+        ray.actor.exit_actor()
+
+    def _build_result(self) -> Result:
+        storage = self._checkpoint_manager._storage_context
+
+        latest_checkpoint_result = self._checkpoint_manager.latest_checkpoint_result
+        latest_metrics = (
+            latest_checkpoint_result.metrics if latest_checkpoint_result else None
+        )
+        latest_checkpoint = (
+            latest_checkpoint_result.checkpoint if latest_checkpoint_result else None
+        )
+        best_checkpoints = [
+            (r.checkpoint, r.metrics)
+            for r in self._checkpoint_manager.best_checkpoint_results
+        ]
+
+        # Provide the history of metrics attached to checkpoints as a dataframe.
+        metrics_dataframe = None
+        if best_checkpoints:
+            metrics_dataframe = pd.DataFrame([m for _, m in best_checkpoints])
+
+        return Result(
+            metrics=latest_metrics,
+            checkpoint=latest_checkpoint,
+            error=self.get_training_failed_error(),
+            path=storage.experiment_fs_path,
+            best_checkpoints=best_checkpoints,
+            metrics_dataframe=metrics_dataframe,
+            _storage_filesystem=storage.storage_filesystem,
+            return_value=self._return_value,
+        )
+
+    def get_result(self) -> Result:
+        """Get the final training result from the TrainController."""
+
+        controller_state = self.get_state()
+        if not controller_state.is_terminal():
+            raise ValueError(
+                f"Cannot get result when controller is in state {controller_state}"
+            )
+        return self._build_result()
+
+    def get_training_failed_error(self) -> Optional[TrainingFailedError]:
+        """Get the training failed error from the controller state.
+
+        Returns:
+            The training failed error if the controller is in an errored state,
+            None otherwise.
+        """
+        controller_state = self.get_state()
+
+        if isinstance(controller_state, ErroredState):
+            return controller_state.training_failed_error
+
+        return None
+
+    async def get_all_reported_checkpoints(
+        self,
+        current_report_index: int,
+        consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
+        timeout_s: Optional[float] = None,
+    ) -> List["ReportedCheckpoint"]:
+        return await self._checkpoint_manager.get_all_reported_checkpoints(
+            current_report_index, consistency_mode, timeout_s
+        )

python/ray/train/v2/_internal/execution/controller/placement_group_cleaner.py

@@ -0,0 +1,183 @@
+import logging
+import queue
+import threading
+from typing import Optional
+
+import ray
+from ray.train.v2._internal.state.util import is_actor_alive
+from ray.util.placement_group import PlacementGroup, remove_placement_group
+
+logger = logging.getLogger(__name__)
+
+
+class PlacementGroupCleaner:
+    """Detached helper that ensures PG cleanup if Ray Train Controller exits ungracefully.
+
+    This actor should be created with lifetime='detached' to avoid being
+    fate-shared with the Train controller.
+    """
+
+    def __init__(
+        self,
+        controller_actor_id: str,
+        check_interval_s: float,
+        get_actor_timeout_s: float,
+        stop_timeout: Optional[float],
+    ):
+        self._controller_actor_id = controller_actor_id
+        self._check_interval_s = check_interval_s
+        self._get_actor_timeout_s = get_actor_timeout_s
+        self._stop_timeout = stop_timeout
+        self._pg_queue: queue.Queue = queue.Queue()
+        self._stop_event = threading.Event()
+        self._monitor_thread: Optional[threading.Thread] = None
+        self._exiting: bool = False
+
+    def register_placement_group(self, placement_group: PlacementGroup):
+        logger.debug(
+            "PlacementGroupCleaner registered placement group %s for controller %s",
+            placement_group.id,
+            self._controller_actor_id,
+        )
+        # Send placement group update to the monitor thread via queue
+        self._pg_queue.put(placement_group)
+
+    def start_monitoring(self):
+        """Start monitoring the controller and placement group."""
+        if self._monitor_thread is not None and self._monitor_thread.is_alive():
+            # Thread already running, just return True
+            logger.debug("Monitor thread already running")
+            return True
+
+        self._monitor_thread = threading.Thread(
+            target=self._monitor_loop,
+            name="PlacementGroupCleanerMonitor",
+            daemon=True,
+        )
+        self._monitor_thread.start()
+        logger.debug("PlacementGroupCleaner started monitoring in background thread")
+        return True
+
+    def _monitor_loop(self):
+        """Monitor controller; remove PG when controller is gone.
+
+        This runs continuously until controller dies or stop() is called.
+        Uses a queue to receive placement group updates.
+        """
+        curr_placement_group: Optional[PlacementGroup] = None
+
+        while not self._stop_event.is_set():
+            # Check for new placement group updates from queue
+            try:
+                pg = self._pg_queue.get(timeout=self._check_interval_s)
+                curr_placement_group = pg
+                logger.debug(f"Updated current placement group to {pg.id}")
+            except queue.Empty:
+                pass  # continue to monitor current placement group
+
+            # Check if controller is still alive
+            try:
+                alive = is_actor_alive(
+                    actor_id=self._controller_actor_id,
+                    timeout=self._get_actor_timeout_s,
+                )
+            except ray.util.state.exception.RayStateApiException:
+                logger.warning(
+                    "Failed to query Ray Train Controller actor state. "
+                    "State API may be temporarily unavailable. Continuing to monitor."
+                )
+                continue
+
+            # Cleanup if controller is dead
+            if not alive:
+                # Drain any queued placement groups
+                while True:
+                    try:
+                        pg = self._pg_queue.get_nowait()
+                        curr_placement_group = pg
+                    except queue.Empty:
+                        break
+                self._cleanup_placement_group(curr_placement_group)
+                break
+
+        # Exit the actor after cleanup since controller is dead
+        self._exit()
+        self._monitor_thread = None
+
+    def _cleanup_placement_group(self, placement_group: Optional[PlacementGroup]):
+        """Clean up the current placement group if it hasn't been removed."""
+        if placement_group is None:
+            logger.debug("No placement group registered; skipping cleanup.")
+            return
+
+        if self._is_placement_group_removed(placement_group):
+            logger.debug(
+                "Controller actor died but placement group already removed; "
+                "skipping cleanup."
+            )
+            return
+
+        logger.warning(
+            f"Detected that the Ray Train controller actor ({self._controller_actor_id}) is dead. "
+            f"Cleaning up placement group = [{placement_group.id}] created by this run."
+        )
+        try:
+            remove_placement_group(placement_group)
+        except Exception as e:
+            logger.warning(f"Failed to clean up placement group: {e}")
+            return
+
+        logger.debug(
+            f"Placement group = [{placement_group.id}] cleaned up successfully"
+        )
+
+    def _stop_monitor_thread(self):
+        """Stop the monitor thread and wait for it to exit.
+
+        Returns:
+            bool: True if the thread was stopped, False if there was no active thread.
+        """
+        if self._monitor_thread is None or not self._monitor_thread.is_alive():
+            return False
+
+        # Signal stop and wait for thread to exit
+        self._stop_event.set()
+        self._monitor_thread.join(timeout=self._stop_timeout)
+        if self._monitor_thread.is_alive():
+            logger.warning(
+                "Monitor thread did not exit within %.2f seconds", self._stop_timeout
+            )
+            return False
+
+        self._monitor_thread = None
+        return True
+
+    def stop(self):
+        """Request the cleaner to stop monitoring and exit."""
+        self._stop_monitor_thread()
+        self._exit()
+
+    def _is_placement_group_removed(self, placement_group: PlacementGroup) -> bool:
+        """Check if a placement group has been removed."""
+        try:
+            table = ray.util.placement_group_table(placement_group)
+        except Exception as e:
+            logger.warning(
+                f"Failed to query placement group table: {e}. "
+                "Assuming placement group is not removed."
+            )
+            return False
+        if "state" not in table:
+            return True
+        return table["state"] == "REMOVED"
+
+    def _exit(self):
+        """Exit the actor."""
+        if self._exiting:
+            return
+        self._exiting = True
+        try:
+            ray.actor.exit_actor()
+        except Exception as e:
+            # If exit fails for any reason, just log it.
+            logger.warning(f"Failed to exit actor: {e}")

python/ray/train/v2/_internal/execution/controller/state.py

@@ -0,0 +1,156 @@
+from enum import Enum
+
+from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
+    ScalingDecision,
+)
+from ray.train.v2.api.exceptions import TrainingFailedError
+
+
+class TrainControllerStateType(Enum):
+    """Enum representing different states of the train controller.
+
+    States:
+       INITIALIZING: The train controller is starting up. This is always the initial
+           state of the controller.
+       SCHEDULING: The train controller is in the process of scheduling a new worker
+           group.
+       RESCHEDULING: The train controller is in the process of rescheduling the worker
+           group.
+       RUNNING: The train controller is actively running training tasks.
+       RESTARTING: The train controller is in the process of recovering from an error.
+       RESIZING: The train controller is in the process of resizing a running worker
+           group.
+       SHUTTING_DOWN: The train controller has already shut down the worker group and
+           and is in the process of shutting itself down.
+       ERRORED: A terminal state indicating that training has encountered an error and
+           cannot continue.
+       FINISHED: A terminal state indicating that training has completed.
+       ABORTED: A terminal state indicating that training has been aborted.
+
+    Args:
+        state_name: The name of the state.
+        is_terminal: Whether this is a terminal state that should not be further processed.
+        needs_new_run_attempt: Whether this state requires starting a new run attempt, where
+            a run attempt is a logical unit that encompasses both scheduling workers and
+            executing training on those workers.
+    """
+
+    INITIALIZING = ("INITIALIZING", False, True)
+    SCHEDULING = ("SCHEDULING", False, False)
+    RESCHEDULING = ("RESCHEDULING", False, False)
+    RUNNING = ("RUNNING", False, False)
+    RESTARTING = ("RESTARTING", False, True)
+    RESIZING = ("RESIZING", False, True)
+    SHUTTING_DOWN = ("SHUTTING_DOWN", False, False)
+    ERRORED = ("ERRORED", True, False)
+    FINISHED = ("FINISHED", True, False)
+    ABORTED = ("ABORTED", True, False)
+
+    def __init__(
+        self,
+        state_name: str,
+        is_terminal: bool,
+        needs_new_run_attempt: bool,
+    ):
+        self.state_name = state_name
+        self.is_terminal = is_terminal
+        self.needs_new_run_attempt = needs_new_run_attempt
+
+
+class TrainControllerState:
+    """Base class for all train controller states.
+
+    Methods:
+        get_type() -> TrainControllerStateType: Returns the type of the state.
+        is_terminal() -> bool: Returns whether the state is terminal.
+        needs_new_run_attempt() -> bool: Returns whether a new run attempt is needed.
+    """
+
+    def __init__(self, state_type: TrainControllerStateType):
+        self._state_type = state_type
+
+    def __repr__(self) -> str:
+        attrs = {
+            "type": self._state_type.name,
+            "is_terminal": self._state_type.is_terminal,
+            "needs_new_run_attempt": self._state_type.needs_new_run_attempt,
+            **{k: v for k, v in vars(self).items() if not k.startswith("_")},
+        }
+        attrs_str = "\n    ".join(f"{k}={v}" for k, v in attrs.items())
+        return f"{self.__class__.__name__}(\n    {attrs_str}\n)"
+
+    def is_terminal(self) -> bool:
+        return self._state_type.is_terminal
+
+    def needs_new_run_attempt(self) -> bool:
+        return self._state_type.needs_new_run_attempt
+
+
+class InitializingState(TrainControllerState):
+    def __init__(self):
+        super().__init__(state_type=TrainControllerStateType.INITIALIZING)
+
+
+class SchedulingState(TrainControllerState):
+    def __init__(self, scaling_decision: ScalingDecision):
+        super().__init__(state_type=TrainControllerStateType.SCHEDULING)
+        self.scaling_decision = scaling_decision
+
+
+class ReschedulingState(TrainControllerState):
+    def __init__(
+        self,
+        training_failed_error: TrainingFailedError,
+    ):
+        super().__init__(state_type=TrainControllerStateType.RESCHEDULING)
+        self.training_failed_error = training_failed_error
+
+
+class RunningState(TrainControllerState):
+    # TODO: Split into multiple more granular states, or add more fields.
+    # For example, we may want to indicate if any health checks failed.
+    def __init__(self):
+        super().__init__(state_type=TrainControllerStateType.RUNNING)
+
+
+class RestartingState(TrainControllerState):
+    def __init__(
+        self,
+        training_failed_error: TrainingFailedError,
+    ):
+        super().__init__(state_type=TrainControllerStateType.RESTARTING)
+        self.training_failed_error = training_failed_error
+
+
+class ResizingState(TrainControllerState):
+    def __init__(
+        self,
+        scaling_decision: ScalingDecision,
+    ):
+        super().__init__(state_type=TrainControllerStateType.RESIZING)
+        self.scaling_decision = scaling_decision
+
+
+class ShuttingDownState(TrainControllerState):
+    def __init__(self, next_state: "TrainControllerState"):
+        super().__init__(state_type=TrainControllerStateType.SHUTTING_DOWN)
+        self.next_state = next_state
+
+
+class ErroredState(TrainControllerState):
+    def __init__(
+        self,
+        training_failed_error: TrainingFailedError,
+    ):
+        super().__init__(state_type=TrainControllerStateType.ERRORED)
+        self.training_failed_error = training_failed_error
+
+
+class FinishedState(TrainControllerState):
+    def __init__(self):
+        super().__init__(state_type=TrainControllerStateType.FINISHED)
+
+
+class AbortedState(TrainControllerState):
+    def __init__(self):
+        super().__init__(state_type=TrainControllerStateType.ABORTED)

python/ray/train/v2/_internal/execution/failure_handling/__init__.py

@@ -0,0 +1,16 @@
+# isort: off
+from .failure_policy import FailureDecision, FailurePolicy
+from .default import DefaultFailurePolicy
+from .factory import create_failure_policy
+
+# isort: on
+
+__all__ = [
+    "DefaultFailurePolicy",
+    "FailureDecision",
+    "FailurePolicy",
+    "create_failure_policy",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/v2/_internal/execution/failure_handling/default.py

@@ -0,0 +1,100 @@
+import logging
+
+from .failure_policy import FailureDecision, FailurePolicy
+from ray.train.v2._internal.exceptions import (
+    WorkerGroupStartupFailedError,
+    WorkerGroupStartupTimeoutError,
+)
+from ray.train.v2.api.config import FailureConfig
+from ray.train.v2.api.exceptions import (
+    ControllerError,
+    TrainingFailedError,
+    WorkerGroupError,
+)
+
+logger = logging.getLogger(__name__)
+
+
+RETRYABLE_CONTROLLER_ERRORS = (
+    WorkerGroupStartupFailedError,
+    WorkerGroupStartupTimeoutError,
+)
+
+
+class DefaultFailurePolicy(FailurePolicy):
+    def __init__(self, failure_config: FailureConfig):
+        super().__init__(failure_config)
+        self._worker_group_failures = 0
+        self._controller_failures = 0
+
+    def _log_decision(
+        self,
+        decision: FailureDecision,
+        training_failed_error: TrainingFailedError,
+        error_count: int,
+        retry_limit: int,
+    ):
+        if isinstance(training_failed_error, ControllerError):
+            error_source = "controller"
+        elif isinstance(training_failed_error, WorkerGroupError):
+            error_source = "worker group"
+        else:
+            raise ValueError(f"Unknown error type: {type(training_failed_error)}")
+
+        logger.info(
+            f"[FailurePolicy] {decision.value}\n"
+            f"  Source: {error_source}\n"
+            f"  Error count: {error_count} (max allowed: {retry_limit})\n"
+            f"Error: {training_failed_error}",
+            exc_info=(
+                type(training_failed_error),
+                training_failed_error,
+                training_failed_error.__traceback__,
+            ),
+        )
+
+    def _is_retryable_error(self, training_failed_error: TrainingFailedError) -> bool:
+        if isinstance(training_failed_error, WorkerGroupError):
+            return True
+        elif isinstance(training_failed_error, ControllerError):
+            return isinstance(
+                training_failed_error.controller_failure, RETRYABLE_CONTROLLER_ERRORS
+            )
+        return False
+
+    def make_decision(
+        self,
+        training_failed_error: TrainingFailedError,
+    ) -> FailureDecision:
+
+        if not self._is_retryable_error(training_failed_error):
+            decision = FailureDecision.RAISE
+            error_count = 1
+            retry_limit = 0
+        else:
+            if isinstance(training_failed_error, ControllerError):
+                self._controller_failures += 1
+                error_count = self._controller_failures
+                retry_limit = (
+                    self.failure_config.controller_failure_limit
+                    if self.failure_config.controller_failure_limit != -1
+                    else float("inf")
+                )
+            elif isinstance(training_failed_error, WorkerGroupError):
+                self._worker_group_failures += 1
+                error_count = self._worker_group_failures
+                retry_limit = (
+                    self.failure_config.max_failures
+                    if self.failure_config.max_failures != -1
+                    else float("inf")
+                )
+            else:
+                raise ValueError(f"Unknown error type: {type(training_failed_error)}")
+
+            if error_count > retry_limit:
+                decision = FailureDecision.RAISE
+            else:
+                decision = FailureDecision.RETRY
+
+        self._log_decision(decision, training_failed_error, error_count, retry_limit)
+        return decision

python/ray/train/v2/_internal/execution/failure_handling/factory.py

@@ -0,0 +1,13 @@
+from ray.train import FailureConfig
+from ray.train.v2._internal.execution.failure_handling import (
+    DefaultFailurePolicy,
+    FailurePolicy,
+)
+
+
+def create_failure_policy(failure_config: FailureConfig) -> FailurePolicy:
+    """Create a failure policy from the given failure config.
+
+    Defaults to the `DefaultFailurePolicy` implementation.
+    """
+    return DefaultFailurePolicy(failure_config=failure_config)

python/ray/train/v2/_internal/execution/failure_handling/failure_policy.py

@@ -0,0 +1,29 @@
+import abc
+from enum import Enum
+
+from ray.train.v2.api.config import FailureConfig
+from ray.train.v2.api.exceptions import TrainingFailedError
+
+
+class FailureDecision(Enum):
+    RETRY = "RETRY"
+    RAISE = "RAISE"
+    NOOP = "NOOP"
+
+
+class FailurePolicy(abc.ABC):
+    """A policy that determines how to handle user and system failures.
+    FailurePolicy will handle the controller failure and worker errors during training.
+
+    This can be used to implement fault tolerance and error recovery.
+    """
+
+    def __init__(self, failure_config: FailureConfig):
+        self.failure_config = failure_config
+
+    @abc.abstractmethod
+    def make_decision(
+        self,
+        training_failed_error: TrainingFailedError,
+    ) -> FailureDecision:
+        raise NotImplementedError

python/ray/train/v2/_internal/execution/local_mode/torch.py

@@ -0,0 +1,93 @@
+import logging
+import os
+from typing import Any, Callable
+
+import torch
+import torch.distributed as dist
+
+from ray.train import Result
+from ray.train.v2._internal.execution.local_mode.utils import LocalController
+from ray.train.v2._internal.execution.train_fn_utils import (
+    LocalTrainFnUtils,
+    get_train_fn_utils,
+    set_train_fn_utils,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def has_torchrun_env() -> bool:
+    """Return True if this process has torch.distributed env vars set.
+
+    For torch.distributed.init_process_group with init_method="env://", these variables are required:
+    - RANK: The rank of the current process
+    - LOCAL_RANK: The local rank of the current process
+    - WORLD_SIZE: Total number of processes participating in the job
+    - LOCAL_WORLD_SIZE: Total number of processes participating in the job on the current node
+    - MASTER_ADDR: The IP address or hostname of the master node (rank 0)
+    - MASTER_PORT: A free port on the master node for communication
+
+    """
+    torch_dist_required_vars = {
+        "RANK",
+        "LOCAL_RANK",
+        "WORLD_SIZE",
+        "LOCAL_WORLD_SIZE",
+        "MASTER_ADDR",
+        "MASTER_PORT",
+    }
+
+    return torch_dist_required_vars.issubset(os.environ.keys())
+
+
+class LocalTorchController(LocalController):
+    def _set_train_fn_utils(self) -> None:
+        world_size = 1
+        global_rank = 0
+        local_rank = 0
+        nproc_per_node = 1
+        node_rank = 0
+        if has_torchrun_env():
+            assert not dist.is_initialized(), "torch.distributed is already initialized"
+            torch.distributed.init_process_group(
+                backend="nccl" if torch.cuda.is_available() else "gloo"
+            )
+            world_size = torch.distributed.get_world_size()
+            global_rank = torch.distributed.get_rank()
+            local_rank = int(os.environ["LOCAL_RANK"])
+            if torch.cuda.is_available():
+                torch.cuda.set_device(local_rank)
+            nproc_per_node = int(os.environ.get("LOCAL_WORLD_SIZE"))
+            node_rank = global_rank // nproc_per_node
+
+        if world_size != 1:
+            assert (
+                self.datasets is None or len(self.datasets) == 0
+            ), "Ray Data is not supported in local mode with multiple workers."
+        set_train_fn_utils(
+            LocalTrainFnUtils(
+                experiment_name=self.experiment_name,
+                world_size=world_size,
+                world_rank=global_rank,
+                local_rank=local_rank,
+                local_world_size=nproc_per_node,
+                node_rank=node_rank,
+                dataset_shards=self.datasets,
+            )
+        )
+
+    def run(self, train_func: Callable[[], Any]) -> Result:
+        self._set_train_fn_utils()
+        train_result = train_func()
+        train_fn_utils = get_train_fn_utils()
+        assert isinstance(train_fn_utils, LocalTrainFnUtils)
+        result = Result(
+            metrics=train_fn_utils._get_last_metrics(),
+            checkpoint=train_fn_utils.get_checkpoint(),
+            path=None,
+            error=None,
+            return_value=train_result,
+        )
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        return result

python/ray/train/v2/_internal/execution/local_mode/utils.py

@@ -0,0 +1,41 @@
+import logging
+from typing import Any, Callable, Dict, Optional
+
+from ray.train import Result
+from ray.train.trainer import GenDataset
+from ray.train.v2._internal.execution.train_fn_utils import (
+    LocalTrainFnUtils,
+    get_train_fn_utils,
+    set_train_fn_utils,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class LocalController:
+    def __init__(
+        self, experiment_name: str, datasets: Optional[Dict[str, GenDataset]] = None
+    ):
+        if datasets is not None:
+            datasets = {k: v() if callable(v) else v for k, v in datasets.items()}
+
+        self.datasets = datasets
+        self.experiment_name = experiment_name
+
+    def run(self, train_func: Callable[[], Any]) -> Result:
+        set_train_fn_utils(
+            LocalTrainFnUtils(
+                experiment_name=self.experiment_name,
+                dataset_shards=self.datasets,
+            )
+        )
+        result = train_func()
+        train_fn_utils = get_train_fn_utils()
+        assert isinstance(train_fn_utils, LocalTrainFnUtils)
+        return Result(
+            metrics=train_fn_utils._get_last_metrics(),
+            checkpoint=train_fn_utils.get_checkpoint(),
+            path=None,
+            error=None,
+            return_value=result,
+        )

python/ray/train/v2/_internal/execution/preemption.py

@@ -0,0 +1,243 @@
+import logging
+import threading
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, List, Optional, Set
+
+import ray
+from ray.actor import ActorHandle
+from ray.train.v2._internal.constants import DEFAULT_PREEMPTION_POLL_INTERVAL_S
+from ray.util.tpu import get_tpu_slice_name_from_node
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.worker import RayTrainWorker
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class PreemptionInfo:
+    """Information about an imminent preemption event.
+
+    Attributes:
+        deadline_ms: Earliest preemption deadline (UNIX time in milliseconds)
+            across all preempted nodes. ``None`` if no deadline was reported.
+        preempted_node_to_ranks: Map of preempted ``node_id`` to the worker ``world_rank``s affected when that node
+            is preempted.
+    """
+
+    deadline_ms: Optional[int]
+    preempted_node_to_ranks: Dict[str, List[int]]
+
+    @property
+    def preempted_node_ids(self) -> List[str]:
+        """Preempted node IDs, sorted lexicographically."""
+        return sorted(self.preempted_node_to_ranks)
+
+    @property
+    def preempted_ranks(self) -> List[int]:
+        """All affected ranks across the preempted nodes, sorted ascending."""
+        return sorted(
+            {r for ranks in self.preempted_node_to_ranks.values() for r in ranks}
+        )
+
+
+@dataclass
+class PreemptionContext:
+    """Thread-shared preemption signal for one worker actor."""
+
+    _preemption_info: Optional[PreemptionInfo] = field(default=None, init=False)
+    _lock: threading.Lock = field(default_factory=threading.Lock, init=False)
+
+    def set(self, info: PreemptionInfo) -> None:
+        with self._lock:
+            self._preemption_info = info
+
+    def get(self) -> Optional[PreemptionInfo]:
+        """Return the current preemption signal, or ``None`` if none received."""
+        with self._lock:
+            return self._preemption_info
+
+
+def _get_draining_nodes() -> Dict[str, int]:
+    """Ray Core's draining nodes as ``{node_id_hex: deadline_ms}`` (0 = no deadline)."""
+    return ray._private.state.state.get_draining_nodes()
+
+
+class PreemptionWatcher:
+    """Polls Ray Core for node drains and logs detected preemption events.
+
+    One watcher per worker group, spawned as a ``num_cpus=0`` actor by
+    ``PreemptionCallback``. The poll loop runs in a background thread. The
+    failure-domain map is built once on construction and is immutable for the
+    watcher's lifetime.
+
+    The failure-domain map records which of our ranks are affected if a node is
+    preempted: for a GPU node, the ranks on that node; for a TPU node, every
+    rank in the node's slice, since a TPU slice is preempted atomically.
+
+    Args:
+        node_to_ranks: Map ``node_id_hex -> [ranks on that node]``. Used both
+            as the set of nodes we care about (drains elsewhere are ignored)
+            and as the seed for failure-domain expansion.
+        poll_interval_s: Seconds between drain-state polls.
+        worker_actors_by_rank: Map ``world_rank -> worker actor handle``. On a
+            detected preemption, ``mark_preempt`` is called on every worker.
+    """
+
+    def __init__(
+        self,
+        node_to_ranks: Dict[str, List[int]],
+        poll_interval_s: float = DEFAULT_PREEMPTION_POLL_INTERVAL_S,
+        worker_actors_by_rank: Optional[
+            Dict[int, ActorHandle["RayTrainWorker"]]
+        ] = None,
+    ):
+        self._node_to_ranks: Dict[str, List[int]] = {
+            nid: sorted(ranks) for nid, ranks in node_to_ranks.items()
+        }
+        self._poll_interval_s = poll_interval_s
+        self._worker_actors_by_rank: Dict[int, ActorHandle["RayTrainWorker"]] = (
+            worker_actors_by_rank or {}
+        )
+        self._failure_domain_map: Dict[str, List[int]] = self._build_failure_domain_map(
+            self._node_to_ranks
+        )
+
+        self._stop_event = threading.Event()
+        self._last_drained: Dict[str, int] = {}
+        self._latest_info: Optional[PreemptionInfo] = None
+
+        self._monitor_thread = threading.Thread(
+            target=self._watch_loop,
+            name="PreemptionWatcher",
+            daemon=True,
+        )
+        self._monitor_thread.start()
+
+    @staticmethod
+    def _build_failure_domain_map(
+        node_to_ranks: Dict[str, List[int]],
+    ) -> Dict[str, List[int]]:
+        """Map each node we host to all ranks in its failure domain.
+
+        - Non-TPU (e.g. GPU) clusters: the failure domain is the node itself,
+          so a drain on a node flags only the ranks this job runs there.
+        - TPU multislice: every host in a slice is reclaimed atomically, so a
+          drain on any host is fate-shared with the rest.
+        """
+        per_node = {nid: sorted(set(ranks)) for nid, ranks in node_to_ranks.items()}
+
+        try:
+            all_nodes = ray.nodes()
+
+            # Slice label for each node we host (None for non-TPU nodes).
+            node_to_slice: Dict[str, Optional[str]] = {
+                node["NodeID"]: get_tpu_slice_name_from_node(node)
+                for node in all_nodes
+                if node["NodeID"] in node_to_ranks
+            }
+
+            # Union our ranks per slice.
+            slice_to_ranks: Dict[str, Set[int]] = {}
+            for node_id, ranks in node_to_ranks.items():
+                slice_label = node_to_slice.get(node_id)
+                if slice_label:
+                    slice_to_ranks.setdefault(slice_label, set()).update(ranks)
+
+            # Non-TPU cluster (or none of our nodes are on a slice): per-node.
+            if not slice_to_ranks:
+                return per_node
+
+            result: Dict[str, List[int]] = {}
+            for node_id, ranks in node_to_ranks.items():
+                slice_label = node_to_slice.get(node_id)
+                if slice_label:
+                    result[node_id] = sorted(slice_to_ranks[slice_label])
+                else:
+                    result[node_id] = sorted(set(ranks))
+            return result
+        except Exception:
+            logger.debug(
+                "Could not build failure-domain map; falling back to per-node "
+                "domains (no TPU-slice expansion).",
+                exc_info=True,
+            )
+            return per_node
+
+    def get_latest_preemption_info(self) -> Optional[PreemptionInfo]:
+        """Most recent :class:`PreemptionInfo` observed, or ``None``."""
+        return self._latest_info
+
+    def _watch_loop(self) -> None:
+        logger.debug(
+            "PreemptionWatcher polling %d node(s) every %.1fs.",
+            len(self._node_to_ranks),
+            self._poll_interval_s,
+        )
+        while not self._stop_event.is_set():
+            self._poll_once()
+            self._stop_event.wait(timeout=self._poll_interval_s)
+        logger.debug("PreemptionWatcher stopped.")
+
+    def _poll_once(self) -> None:
+        """Poll the drain source once and dispatch on change.
+
+        Per-poll exceptions are caught and logged so a transient GCS hiccup
+        doesn't kill the watcher loop.
+        """
+        try:
+            drained = _get_draining_nodes() or {}
+            # Keep only drains on this job's own nodes (others are ignored).
+            # That's complete for TPU — an SPMD job fully occupies its slice, so
+            # every fate-shared host is one of our nodes and a drain on any slice
+            # host appears here. For GPU, a drain on a host we don't run on is
+            # correctly irrelevant.
+            relevant = {
+                n: d for n, d in drained.items() if n in self._failure_domain_map
+            }
+            if relevant != self._last_drained:
+                self._on_drain_change(relevant)
+                self._last_drained = relevant
+        except Exception:
+            # TODO(lehui): consider exponential backoff when the drain API keeps
+            # failing, instead of retrying at the fixed poll interval.
+            logger.warning("PreemptionWatcher poll failed", exc_info=True)
+
+    def _on_drain_change(self, drained: Dict[str, int]) -> None:
+        """Handle a change in the drained-node set.
+
+        ``drained`` has already been narrowed to this job's nodes by the
+        caller (``_poll_once``).
+        """
+        if not drained:
+            return
+
+        affected_node_ids = sorted(drained.keys())
+        preempted_node_to_ranks = {
+            node_id: self._failure_domain_map[node_id] for node_id in affected_node_ids
+        }
+
+        # Earliest deadline across the preempted nodes; None if none reported one
+        # (Ray Core uses 0 for "no deadline", which is falsy and filtered out).
+        reported_deadlines = [drained[n] for n in affected_node_ids if drained[n]]
+        deadline_ms = min(reported_deadlines) if reported_deadlines else None
+
+        info = PreemptionInfo(
+            deadline_ms=deadline_ms,
+            preempted_node_to_ranks=preempted_node_to_ranks,
+        )
+        self._latest_info = info
+
+        logger.warning(
+            "PreemptionWatcher: preemption detected — "
+            "preempted_node_ids=%s, preempted_ranks=%s, deadline_ms=%s",
+            info.preempted_node_ids,
+            info.preempted_ranks,
+            deadline_ms,
+        )
+
+        for rank, actor in self._worker_actors_by_rank.items():
+            actor.mark_preempt.remote(info)
+        # TODO(lehui): coalesce preemptions seen within one window into a single
+        # worker-group restart, so a staggered drain (node A at t, node B at
+        # t+60s) doesn't cause back-to-back restarts.

python/ray/train/v2/_internal/execution/scaling_policy/__init__.py

@@ -0,0 +1,29 @@
+# isort: off
+from .scaling_policy import ScalingDecision, ScalingPolicy, NoopDecision, ResizeDecision
+from .scaling_policy import (
+    AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
+    AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+    AUTOSCALING_REQUESTS_INTERVAL_S,
+)
+from .elastic import ElasticScalingPolicy
+from .fixed import FixedScalingPolicy
+from .factory import create_scaling_policy
+
+# isort: on
+
+
+__all__ = [
+    "AUTOSCALING_REQUESTS_EXPIRE_TIME_S",
+    "AUTOSCALING_REQUESTS_GET_TIMEOUT_S",
+    "AUTOSCALING_REQUESTS_INTERVAL_S",
+    "ScalingPolicy",
+    "ElasticScalingPolicy",
+    "FixedScalingPolicy",
+    "ScalingDecision",
+    "NoopDecision",
+    "ResizeDecision",
+    "create_scaling_policy",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/v2/_internal/execution/scaling_policy/elastic.py

@@ -0,0 +1,291 @@
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import ray
+from ray.train.v2._internal.execution.scaling_policy import (
+    NoopDecision,
+    ResizeDecision,
+    ScalingDecision,
+    ScalingPolicy,
+)
+from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
+    AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+)
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroupPollStatus,
+    WorkerGroupState,
+)
+from ray.train.v2._internal.util import time_monotonic
+from ray.train.v2.api.config import ScalingConfig
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
+        ResourceDict,
+    )
+
+
+class ElasticScalingPolicy(ScalingPolicy):
+
+    # Minimum interval in seconds between querying the AutoscalingCoordinator for allocated resources.
+    GET_ALLOCATED_RESOURCES_INTERVAL_S = 1
+    # Minimum interval in seconds between logging warnings about insufficient workers.
+    INSUFFICIENT_WORKERS_WARNING_INTERVAL_S = 30
+
+    def __init__(self, scaling_config: ScalingConfig):
+        super().__init__(scaling_config)
+
+        self._latest_monitor_time = float("-inf")
+        self._latest_insufficient_workers_warning_time = float("-inf")
+        self._latest_allocated_resources_query_time = float("-inf")
+        self._latest_allocated_resources: Optional[List["ResourceDict"]] = None
+
+    def _get_num_workers_for_resource_request(self) -> int:
+        return self.scaling_config.max_workers
+
+    def _count_possible_workers(
+        self, allocated_resources: List[Dict[str, float]]
+    ) -> int:
+        """Count the number of workers that can be started/restarted with the given
+        the list of node resources. The returned number is capped at the maximum
+        number of workers.
+
+        For GPUs, this divides raw allocated resources by per-worker requirements.
+        For TPUs, an additional check ensures workers align with physically intact
+        TPU slices (see ``_get_strict_tpu_worker_count``).
+
+        Args:
+            allocated_resources: The resources currently allocated by the AutoscalingCoordinator.
+
+        Returns:
+            The number of workers that can be started/restarted with the current resources.
+        """
+        # TODO: Fractional resources do not work well here.
+        single_worker_resources = self.scaling_config._resources_per_worker_not_none
+        total_num_workers = 0
+
+        # If workers require no resources, we can run as many as we want.
+        if sum(single_worker_resources.values()) == 0:
+            return self.scaling_config.max_workers
+
+        for resources in allocated_resources:
+            num_workers = min(
+                [
+                    resources.get(resource, 0.0) // single_worker_resources[resource]
+                    for resource in single_worker_resources
+                    if single_worker_resources[resource] > 0
+                ]
+            )
+            total_num_workers += num_workers
+
+        total_num_workers = min(int(total_num_workers), self.scaling_config.max_workers)
+
+        # Multi-host TPUs are scheduled atomically in interconnected slices defined by a topology.
+        if (
+            self.scaling_config.use_tpu
+            and self.scaling_config.topology
+            and self.scaling_config.accelerator_type
+        ):
+            total_num_workers = self._get_strict_tpu_worker_count(
+                total_num_workers=total_num_workers,
+            )
+
+        return total_num_workers
+
+    def _get_strict_tpu_worker_count(self, total_num_workers: int) -> int:
+        """Calculate the number of workers that can run on intact TPU slices.
+
+        The Autoscaler's allocated resources might overestimate the number of
+        schedulable TPU workers because it counts raw resources. TPUs require
+        atomic, interconnected slices. This function checks the cluster for
+        physically intact slices to prevent scaling onto fractional/broken
+        topologies.
+
+        The worker count is: min(resource_based_slices, intact_slices) *
+        workers_per_slice, where resource_based_slices =
+        total_num_workers // workers_per_slice.
+
+        Args:
+            total_num_workers: The initial estimate of workers based on raw
+                allocated resources.
+
+        Returns:
+            The number of workers aligned to fully intact TPU slices.
+        """
+        from ray.util.tpu import get_num_tpu_slices, get_tpu_worker_resources
+
+        single_worker_resources = self.scaling_config._resources_per_worker_not_none
+
+        try:
+            workers_per_slice, _ = get_tpu_worker_resources(
+                topology=self.scaling_config.topology,
+                accelerator_type=self.scaling_config.accelerator_type,
+                resources_per_unit=single_worker_resources,
+                num_slices=1,
+            )
+
+            if workers_per_slice == 0:
+                # A single worker requires more resources than exist in a
+                # full slice — impossible scheduling configuration for TPU.
+                return 0
+
+            num_slices_from_resources = total_num_workers // workers_per_slice
+
+            if num_slices_from_resources > 0:
+                try:
+                    num_intact_slices = get_num_tpu_slices(
+                        topology=self.scaling_config.topology,
+                        accelerator_type=self.scaling_config.accelerator_type,
+                    )
+                    num_slices_from_resources = min(
+                        num_slices_from_resources, num_intact_slices
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to check cluster state for intact TPU slices: {e}"
+                    )
+
+            return num_slices_from_resources * workers_per_slice
+
+        except Exception as e:
+            logger.warning(
+                f"Could not calculate TPU slice boundaries for elastic scaling: {e}. "
+                "Worker counts may not align with TPU topology."
+            )
+
+        return 0
+
+    def _get_resize_decision(self, num_workers: int) -> ResizeDecision:
+        return ResizeDecision(
+            num_workers=num_workers,
+            resources_per_worker=self.scaling_config._resources_per_worker_not_none,
+        )
+
+    def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
+        self._maybe_send_resource_request()
+
+        allocated_resources = self._get_allocated_resources()
+        if allocated_resources is None:
+            return NoopDecision()
+
+        num_workers = self._count_possible_workers(allocated_resources)
+
+        if num_workers < self.scaling_config.min_workers:
+            now = time_monotonic()
+            # Only log this warning periodically to avoid spamming logs
+            if (
+                now - self._latest_insufficient_workers_warning_time
+                >= self.INSUFFICIENT_WORKERS_WARNING_INTERVAL_S
+            ):
+                logger.info(
+                    f"Detected ready resources for {num_workers} workers "
+                    "in the cluster. "
+                    "Deciding NOT to start/restart training due to the number of workers "
+                    "falling below the minimum "
+                    f"(min_workers={self.scaling_config.min_workers})."
+                )
+                self._latest_insufficient_workers_warning_time = now
+            return NoopDecision()
+
+        logger.info(
+            f"Detected ready resources for {num_workers} workers "
+            "in the cluster. "
+            "Deciding to start/restart training with this worker group size."
+        )
+        return self._get_resize_decision(num_workers)
+
+    def make_decision_for_running_worker_group(
+        self,
+        worker_group_state: WorkerGroupState,
+        worker_group_status: WorkerGroupPollStatus,
+    ) -> ScalingDecision:
+        self._maybe_send_resource_request()
+
+        # Ensure that we don't make resizing decisions too frequently.
+        # The latest restart time and the latest monitor time (whichever is later)
+        # determine the time of the next resize consideration.
+        latest_consideration_time = max(
+            worker_group_state.start_time, self._latest_monitor_time
+        )
+
+        now = time_monotonic()
+        time_since_latest_consideration = now - latest_consideration_time
+        if (
+            time_since_latest_consideration
+            < self.scaling_config.elastic_resize_monitor_interval_s
+        ):
+            logger.debug(
+                "Skipping resize decision due to the latest resizing consideration "
+                "happening too recently: "
+                "%.2f seconds < ScalingConfig(elastic_resize_monitor_interval_s=%.2f seconds).",
+                time_since_latest_consideration,
+                self.scaling_config.elastic_resize_monitor_interval_s,
+            )
+            return NoopDecision()
+
+        self._latest_monitor_time = now
+
+        allocated_resources = self._get_allocated_resources()
+        if allocated_resources is None:
+            return NoopDecision()
+
+        num_workers = self._count_possible_workers(allocated_resources)
+
+        if num_workers == worker_group_state.num_workers:
+            logger.info(
+                "Did not detect any changes in the cluster resources. "
+                "Training will continue with the same worker group size "
+                f"({num_workers})."
+            )
+            return NoopDecision()
+        elif num_workers < self.scaling_config.min_workers:
+            # This covers an edge case where allocated resources decrease to less
+            # than the minimum number of workers.
+            # This situation is rare, since cluster downsizing typically involves
+            # worker failures. However, this check is still useful to fully
+            # avoid entering an invalid state with fewer workers than the minimum.
+            return NoopDecision()
+
+        logger.info(
+            "Detected changes in the cluster resources. "
+            "Deciding to resize the worker group from "
+            f"{worker_group_state.num_workers} -> {num_workers} workers."
+        )
+        return self._get_resize_decision(num_workers)
+
+    # ---------------------------------------------------
+    # Methods for interacting with AutoscalingCoordinator
+    # ---------------------------------------------------
+
+    def _get_allocated_resources(self) -> Optional[List["ResourceDict"]]:
+        """Get allocated resources from AutoscalingCoordinator.
+        Return None if there is an error."""
+        now = time_monotonic()
+        time_since_last_call = now - self._latest_allocated_resources_query_time
+
+        if time_since_last_call < self.GET_ALLOCATED_RESOURCES_INTERVAL_S:
+            return self._latest_allocated_resources
+
+        allocated_resources = None
+        try:
+            allocated_resources = ray.get(
+                self._autoscaling_coordinator.get_allocated_resources.remote(
+                    self._requester_id
+                ),
+                timeout=AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+            )
+        except Exception:
+            msg = (
+                f"Failed to get allocated resources for {self._requester_id}."
+                " Will not resize the worker group."
+                " If this only happens transiently during network partition or"
+                " CPU being overloaded, it's safe to ignore this error."
+                " If this error persists, file a GitHub issue."
+            )
+            logger.warning(msg, exc_info=True)
+        finally:
+            self._latest_allocated_resources_query_time = time_monotonic()
+            self._latest_allocated_resources = allocated_resources
+
+        return self._latest_allocated_resources

python/ray/train/v2/_internal/execution/scaling_policy/factory.py

@@ -0,0 +1,16 @@
+from ray.train.v2._internal.execution.scaling_policy import (
+    ElasticScalingPolicy,
+    FixedScalingPolicy,
+    ScalingPolicy,
+)
+from ray.train.v2.api.config import ScalingConfig
+
+
+def create_scaling_policy(scaling_config: ScalingConfig) -> ScalingPolicy:
+    """Create a scaling policy from the given scaling config.
+
+    Defaults to the `FixedScalingPolicy` implementation.
+    """
+    if scaling_config.elasticity_enabled:
+        return ElasticScalingPolicy(scaling_config=scaling_config)
+    return FixedScalingPolicy(scaling_config=scaling_config)

python/ray/train/v2/_internal/execution/scaling_policy/fixed.py

@@ -0,0 +1,30 @@
+from ray.train.v2._internal.execution.scaling_policy import (
+    NoopDecision,
+    ResizeDecision,
+    ScalingDecision,
+    ScalingPolicy,
+)
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroupPollStatus,
+    WorkerGroupState,
+)
+
+
+class FixedScalingPolicy(ScalingPolicy):
+    def _get_num_workers_for_resource_request(self) -> int:
+        return self.scaling_config.num_workers
+
+    def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
+        self._maybe_send_resource_request()
+        return ResizeDecision(
+            num_workers=self.scaling_config.num_workers,
+            resources_per_worker=self.scaling_config._resources_per_worker_not_none,
+        )
+
+    def make_decision_for_running_worker_group(
+        self,
+        worker_group_state: WorkerGroupState,
+        worker_group_status: WorkerGroupPollStatus,
+    ) -> ScalingDecision:
+        self._maybe_send_resource_request()
+        return NoopDecision()

python/ray/train/v2/_internal/execution/scaling_policy/scaling_policy.py

@@ -0,0 +1,183 @@
+import abc
+import logging
+import uuid
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Dict
+
+import ray
+from ray.train.v2._internal.execution.callback import ControllerCallback
+from ray.train.v2._internal.execution.context import TrainRunContext
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroupPollStatus,
+    WorkerGroupState,
+)
+from ray.train.v2._internal.util import time_monotonic
+from ray.train.v2.api.config import ScalingConfig
+
+logger = logging.getLogger(__name__)
+
+# The time in seconds after which an autoscaling request will expire.
+AUTOSCALING_REQUESTS_EXPIRE_TIME_S = 180
+# Timeout in seconds for getting the result of a call to the AutoscalingCoordinator.
+AUTOSCALING_REQUESTS_GET_TIMEOUT_S = 5
+# Interval in seconds between resource requests to the AutoscalingCoordinator.
+AUTOSCALING_REQUESTS_INTERVAL_S = 20
+
+
+@dataclass
+class ScalingDecision:
+    pass
+
+
+@dataclass
+class NoopDecision(ScalingDecision):
+    pass
+
+
+@dataclass
+class ResizeDecision(ScalingDecision):
+    num_workers: int
+    resources_per_worker: Dict[str, float]
+
+
+class ScalingPolicy(abc.ABC, ControllerCallback):
+    """A policy that determines when and how to scale a worker group.
+
+    This can be used to implement elasticity and fault tolerance.
+
+    Recovery decisions are made when workers are in an inactive or unhealthy state.
+    Upscale decisions are optional and are made when workers are healthy.
+
+    Note: When adding new scaling policies, revisit the shared defaults- particularly if:
+    - AutoscalingCoordinator integration is not needed or a different interface
+      becomes available
+    - Timeout/expiry constants need to diverge between policies
+    - _get_num_workers_for_resource_request() needs variable worker counts
+    - Controller lifecycle behavior diverges
+    """
+
+    # TODO: Restructure these APIs to consider different TrainControllerStates
+    # instead of just running and non-running worker groups.
+
+    def __init__(self, scaling_config: ScalingConfig):
+        self.scaling_config = scaling_config
+        self._requester_id = "train-" + uuid.uuid4().hex
+        self._latest_autoscaling_request_time = float("-inf")
+
+    @abc.abstractmethod
+    def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
+        """Makes a scaling decision when the worker group is initializing
+        or recovering from an error."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def make_decision_for_running_worker_group(
+        self,
+        worker_group_state: WorkerGroupState,
+        worker_group_status: WorkerGroupPollStatus,
+    ) -> ScalingDecision:
+        """Makes a scaling decision when monitoring healthy, running workers."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def _get_num_workers_for_resource_request(self) -> int:
+        """Return the number of workers to request resources for."""
+        raise NotImplementedError
+
+    # ---------------------------------------------------
+    # Methods for interacting with AutoscalingCoordinator
+    # ---------------------------------------------------
+
+    @cached_property
+    def _autoscaling_coordinator(self):
+        from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
+            get_or_create_autoscaling_coordinator,
+        )
+
+        return get_or_create_autoscaling_coordinator()
+
+    def _maybe_send_resource_request(self):
+        """Send a resource request to AutoscalingCoordinator,
+        if AUTOSCALING_REQUESTS_INTERVAL_S has passed since the last send."""
+        now = time_monotonic()
+        if (
+            now - self._latest_autoscaling_request_time
+            < AUTOSCALING_REQUESTS_INTERVAL_S
+        ):
+            return
+        self._send_resource_request()
+
+    def _send_resource_request(self):
+        """Register training resources with the AutoscalingCoordinator."""
+        resources_per_worker = self.scaling_config._resources_per_worker_not_none
+        num_workers = self._get_num_workers_for_resource_request()
+        label_selectors = self.scaling_config._label_selector_per_worker(num_workers)
+        try:
+            from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
+                ResourceRequestPriority,
+            )
+
+            ray.get(
+                self._autoscaling_coordinator.request_resources.remote(
+                    requester_id=self._requester_id,
+                    resources=[resources_per_worker] * num_workers,
+                    label_selectors=label_selectors,
+                    expire_after_s=AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
+                    priority=ResourceRequestPriority.HIGH,
+                ),
+                timeout=AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+            )
+            self._latest_autoscaling_request_time = time_monotonic()
+        except Exception:
+            msg = (
+                f"Failed to send resource request for {self._requester_id}."
+                " If this only happens transiently during network partition or"
+                " CPU being overloaded, it's safe to ignore this error."
+                " If this error persists, file a GitHub issue."
+            )
+            logger.warning(msg, exc_info=True)
+
+    def _cancel_resource_request(self):
+        """Cancel the resource request to AutoscalingCoordinator."""
+        try:
+            ray.get(
+                self._autoscaling_coordinator.cancel_request.remote(
+                    requester_id=self._requester_id,
+                ),
+                timeout=AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
+            )
+        except Exception:
+            msg = (
+                f"Failed to cancel resource request for {self._requester_id}."
+                " The request will still expire after the timeout of"
+                f" {AUTOSCALING_REQUESTS_EXPIRE_TIME_S} seconds."
+            )
+            logger.warning(msg, exc_info=True)
+
+    # --------------------------
+    # ControllerCallback
+    # --------------------------
+
+    def after_controller_start(self, train_run_context: TrainRunContext):
+        """Register training resources with the AutoscalingCoordinator."""
+        self._requester_id = f"train-{train_run_context.run_id}"
+        resources_per_worker = self.scaling_config._resources_per_worker_not_none
+        num_workers = self._get_num_workers_for_resource_request()
+        label_selectors = self.scaling_config._label_selector_per_worker(num_workers)
+        if label_selectors:
+            logger.info(
+                f"Requesting resources: {resources_per_worker} * {num_workers} "
+                f"with label_selectors={label_selectors}"
+            )
+        else:
+            logger.info(f"Requesting resources: {resources_per_worker} * {num_workers}")
+        self._send_resource_request()
+
+    async def before_controller_shutdown(self):
+        """Cancel the resource request when the controller shuts down."""
+        self._cancel_resource_request()
+
+    def before_controller_abort(self):
+        """Cancel the resource request when the controller is aborted."""
+        self._cancel_resource_request()

python/ray/train/v2/_internal/execution/storage.py

@@ -0,0 +1,573 @@
+# 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
+# isort: on
+
+import fnmatch
+import logging
+import os
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Type, Union
+
+from ray.air._internal.filelock import TempFileLock
+from ray.train.constants import _get_ray_train_session_dir
+from ray.train.v2._internal.constants import (
+    CHECKPOINT_MANAGER_SNAPSHOT_FILENAME,
+    VALIDATE_STORAGE_MARKER_FILENAME,
+)
+from ray.train.v2._internal.util import date_str
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.train import Checkpoint
+
+
+logger = logging.getLogger(__name__)
+
+
+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: Additional keyword arguments forwarded to
+            ``pyarrow.fs.LocalFileSystem``.
+    """
+
+    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 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):
+    """Deletes (fs, fs_path) or raises FileNotFoundError if it doesn't exist."""
+    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: Callable[[pyarrow.fs.FileInfo], bool] = lambda x: True,
+) -> 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.
+    """
+    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 query.
+        fs_path: The path on the filesystem.
+
+    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 ``(filesystem, path)`` tuple.
+    """
+    storage_path = str(storage_path)
+
+    if storage_filesystem:
+        return storage_filesystem, storage_path
+
+    return pyarrow.fs.FileSystem.from_uri(storage_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'
+        >>> 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,
+        storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
+        read_only: bool = False,
+    ):
+        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.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.read_only = read_only
+        if not self.read_only:
+            self._create_validation_file()
+            self._check_validation_file()
+
+    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"
+            ">"
+        )
+
+    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 persist_current_checkpoint(
+        self, checkpoint: "Checkpoint", checkpoint_dir_name: str
+    ) -> "Checkpoint":
+        """Persists a given checkpoint to the current checkpoint path on the filesystem.
+
+        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, experiment_fs_path / checkpoint_dir_name).
+            checkpoint_dir_name: Name of the destination directory for the
+                checkpoint, relative to ``experiment_fs_path``.
+
+        Returns:
+            Checkpoint: A Checkpoint pointing to the persisted checkpoint location.
+        """
+        if self.read_only:
+            raise RuntimeError(
+                "Cannot perform write/validation operations as the StorageContext is read-only."
+            )
+
+        # TODO(justinvyu): Fix this cyclical import.
+        from ray.train import Checkpoint
+
+        checkpoint_fs_path = self.build_checkpoint_path_from_name(checkpoint_dir_name)
+
+        logger.debug(
+            "Copying checkpoint files to storage path:\n"
+            "({source_fs}, {source}) -> ({dest_fs}, {destination})".format(
+                source=checkpoint.path,
+                destination=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(checkpoint_fs_path)
+        _pyarrow_fs_copy_files(
+            source=checkpoint.path,
+            destination=checkpoint_fs_path,
+            source_filesystem=checkpoint.filesystem,
+            destination_filesystem=self.storage_filesystem,
+        )
+
+        persisted_checkpoint = Checkpoint(
+            filesystem=self.storage_filesystem,
+            path=checkpoint_fs_path,
+        )
+        logger.info(f"Checkpoint successfully created at: {persisted_checkpoint}")
+        return persisted_checkpoint
+
+    @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()
+
+    @property
+    def local_working_directory(self) -> str:
+        """Every ray train worker will set this directory as its working directory."""
+        if self.experiment_dir_name is None:
+            raise RuntimeError(
+                "Cannot access `local_working_directory` without "
+                "setting `experiment_dir_name`"
+            )
+        return Path(_get_ray_train_session_dir(), self.experiment_dir_name).as_posix()
+
+    @property
+    def checkpoint_manager_snapshot_path(self) -> str:
+        """The path to the checkpoint manager snapshot file."""
+        return Path(
+            self.experiment_fs_path, CHECKPOINT_MANAGER_SNAPSHOT_FILENAME
+        ).as_posix()
+
+    @staticmethod
+    def get_experiment_dir_name(run_obj: Union[str, Callable, Type]) -> str:
+        from ray.tune.experiment import Experiment
+
+        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_default_checkpoint_dir_name():
+        """Get the name of the checkpoint directory by timestamp."""
+        return f"checkpoint_{date_str(include_ms=True)}"
+
+    def extract_checkpoint_dir_name_from_path(self, checkpoint_path: str) -> str:
+        """Get the checkpoint name from the checkpoint path.
+        The parent directory of the checkpoint path should be the experiment directory.
+        """
+        # TODO: Use Pathlib to extract the name when supports at least Python 3.9
+        experiment_fs_path = self.experiment_fs_path + "/"
+        if not checkpoint_path.startswith(experiment_fs_path):
+            raise ValueError(
+                f"Checkpoint path {checkpoint_path} is not under the experiment "
+                f"directory {self.experiment_fs_path}."
+            )
+        return checkpoint_path[len(experiment_fs_path) :]
+
+    def build_checkpoint_path_from_name(self, checkpoint_name: str) -> str:
+        """Get the checkpoint path from the checkpoint name.
+        The parent directory of the checkpoint path should be the experiment directory.
+        """
+        return Path(self.experiment_fs_path, checkpoint_name).as_posix()

python/ray/train/v2/_internal/execution/train_fn_utils.py

@@ -0,0 +1,297 @@
+import logging
+import threading
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+
+from ray.train.v2._internal.data_integration.interfaces import DatasetShardMetadata
+from ray.train.v2._internal.execution import collective_impl
+from ray.train.v2._internal.execution.context import (
+    get_train_context as get_internal_train_context,
+)
+from ray.train.v2.api.context import (
+    DistributedTrainContext,
+    LocalTrainContext,
+    TrainContext as ExternalTrainContext,
+)
+from ray.train.v2.api.report_config import (
+    CheckpointConsistencyMode,
+    CheckpointUploadMode,
+)
+from ray.train.v2.api.validation_config import ValidationTaskConfig
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from ray.data import DataIterator
+    from ray.train import Checkpoint
+    from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
+
+
+class TrainFnUtils(ABC):
+    """Utility class providing an abstraction layer between user-facing APIs
+        and :class:`~ray.train.v2.api.context.TrainContext`.
+
+    It should be set before the users' training function is called.
+    This class can be patched if new user APIs behaviors is wanted.
+    """
+
+    @abstractmethod
+    def report(
+        self,
+        metrics: Dict[str, Any],
+        checkpoint: Optional["Checkpoint"] = None,
+        checkpoint_dir_name: Optional[str] = None,
+        checkpoint_upload_mode: CheckpointUploadMode = CheckpointUploadMode.SYNC,
+        delete_local_checkpoint_after_upload: Optional[bool] = None,
+        checkpoint_upload_fn: Optional[
+            Callable[["Checkpoint", str], "Checkpoint"]
+        ] = None,
+        validation: Union[bool, ValidationTaskConfig] = False,
+    ) -> None:
+        """Upload checkpoint to remote storage and put a training result on the result queue.
+
+        Args:
+            metrics: The metrics to report.
+            checkpoint: The checkpoint to report.
+            checkpoint_dir_name: The name of the checkpoint dir
+                in this iteration. Note: If not set, the checkpoint will
+                be stored in the default storage path. If set, make sure
+                this value is unique for each iteration.
+            checkpoint_upload_mode: The manner in which we want to upload the checkpoint.
+                Defaults to uploading the checkpoint synchronously.
+                This works when no checkpoint is provided but is not useful in that case.
+            delete_local_checkpoint_after_upload: Whether to delete the checkpoint after it is uploaded.
+            checkpoint_upload_fn: A user defined function that will be called with the
+                checkpoint to upload it. If not provided, defaults to using the `pyarrow.fs.copy_files`
+                utility for copying to the destination `storage_path`.
+            validation: [Alpha] If True, triggers validation with default kwargs from validation_config.
+                If a ValidationTaskConfig, validation is run using fn_kwargs merged with validation_config
+                defaults, with fn_kwargs taking precedence on conflicts. If False, no validation.
+        """
+        pass
+
+    @abstractmethod
+    def get_checkpoint(self) -> Optional["Checkpoint"]:
+        """Get the latest checkpoint to resume training from.
+
+        Returns:
+            The latest checkpoint if available, None otherwise.
+        """
+        pass
+
+    @abstractmethod
+    def get_all_reported_checkpoints(
+        self,
+        consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
+        timeout_s: Optional[float] = None,
+    ) -> List["ReportedCheckpoint"]:
+        """Get all the checkpoints reported by the workers.
+
+        Args:
+            consistency_mode: Read semantics for checkpoint retrieval. Defaults to VALIDATED.
+            timeout_s: Timeout in seconds for reading checkpoints and validation data.
+                Defaults to ``None`` to not time out.
+
+        Returns:
+            A list of ReportedCheckpoint objects that represent the checkpoints and
+            corresponding metrics reported by the workers.
+        """
+        pass
+
+    @abstractmethod
+    def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
+        """Get the dataset shard for this training process.
+
+        Args:
+            dataset_info: The metadata of the dataset to get the shard for.
+
+        Returns:
+            The DataIterator shard for this worker.
+        """
+        pass
+
+    @abstractmethod
+    def get_context(self) -> ExternalTrainContext:
+        """Get the TrainContext for this training process.
+        The specific type of TrainContext returned depends on the implementation of TrainFnUtils.
+
+        Returns:
+            The train context for this training process.
+        """
+        pass
+
+    @abstractmethod
+    def is_distributed(self) -> bool:
+        pass
+
+    @abstractmethod
+    def barrier(self) -> None:
+        """Create a barrier across all workers.
+
+        All workers must call this method before the training function can continue.
+
+        This method is used by the public API function :func:`ray.train.collective.barrier`.
+        Users should typically call ``ray.train.collective.barrier()`` instead of calling this method directly.
+        """
+        pass
+
+    @abstractmethod
+    def broadcast_from_rank_zero(self, data: Any) -> Any:
+        """Broadcast data from the rank 0 worker to all other workers.
+
+        This method is used by the public API function :func:`ray.train.collective.broadcast_from_rank_zero`.
+        Users should typically call ``ray.train.collective.broadcast_from_rank_zero()`` instead of calling this method directly.
+        """
+        pass
+
+
+class DistributedTrainFnUtils(TrainFnUtils):
+    def report(
+        self,
+        metrics: Dict[str, Any],
+        checkpoint: Optional["Checkpoint"] = None,
+        checkpoint_dir_name: Optional[str] = None,
+        checkpoint_upload_mode: CheckpointUploadMode = CheckpointUploadMode.SYNC,
+        delete_local_checkpoint_after_upload: Optional[bool] = None,
+        checkpoint_upload_fn: Optional[
+            Callable[["Checkpoint", str], "Checkpoint"]
+        ] = None,
+        validation: Union[bool, ValidationTaskConfig] = False,
+    ) -> None:
+        return get_internal_train_context().report(
+            metrics,
+            checkpoint,
+            checkpoint_dir_name,
+            checkpoint_upload_mode,
+            delete_local_checkpoint_after_upload,
+            checkpoint_upload_fn,
+            validation,
+        )
+
+    def get_checkpoint(self):
+        return get_internal_train_context().get_checkpoint()
+
+    def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
+        return get_internal_train_context().get_dataset_shard(dataset_info)
+
+    def get_context(self) -> DistributedTrainContext:
+        return DistributedTrainContext()
+
+    def is_distributed(self) -> bool:
+        return True
+
+    def barrier(self) -> None:
+        return collective_impl.barrier()
+
+    def broadcast_from_rank_zero(self, data: Any) -> Any:
+        return collective_impl.broadcast_from_rank_zero(data)
+
+    def get_all_reported_checkpoints(
+        self,
+        consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
+        timeout_s: Optional[float] = None,
+    ) -> List["ReportedCheckpoint"]:
+        return get_internal_train_context().get_all_reported_checkpoints(
+            consistency_mode=consistency_mode, timeout_s=timeout_s
+        )
+
+
+class LocalTrainFnUtils(TrainFnUtils):
+    def __init__(
+        self,
+        experiment_name: str,
+        dataset_shards: Optional[Dict[str, "DataIterator"]] = None,
+        world_size: int = 1,
+        world_rank: int = 0,
+        local_rank: int = 0,
+        local_world_size: int = 1,
+        node_rank: int = 0,
+    ):
+        self._context = LocalTrainContext(
+            experiment_name=experiment_name,
+            world_size=world_size,
+            world_rank=world_rank,
+            local_rank=local_rank,
+            local_world_size=local_world_size,
+            node_rank=node_rank,
+        )
+        self._dataset_shards = dataset_shards
+        self._last_metrics = None
+        self._last_checkpoint = None
+
+    def report(
+        self,
+        metrics: Dict[str, Any],
+        checkpoint: Optional["Checkpoint"] = None,
+        checkpoint_dir_name: Optional[str] = None,
+        checkpoint_upload_mode: CheckpointUploadMode = CheckpointUploadMode.SYNC,
+        delete_local_checkpoint_after_upload: Optional[bool] = None,
+        checkpoint_upload_fn: Optional[
+            Callable[["Checkpoint", str], "Checkpoint"]
+        ] = None,
+        validation: Union[bool, ValidationTaskConfig] = False,
+    ) -> None:
+        self._last_metrics = metrics
+        self._last_checkpoint = checkpoint
+        logger.info(f"Reported metrics: {metrics}")
+
+    def get_checkpoint(self) -> Optional["Checkpoint"]:
+        return self._last_checkpoint
+
+    def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
+        dataset_name = dataset_info.dataset_name
+        assert (
+            self._dataset_shards is not None and dataset_name in self._dataset_shards
+        ), f"Dataset shard {dataset_name} not found."
+        return self._dataset_shards[dataset_name]
+
+    def get_context(self) -> LocalTrainContext:
+        return self._context
+
+    def is_distributed(self) -> bool:
+        return False
+
+    def barrier(self) -> None:
+        pass
+
+    def broadcast_from_rank_zero(self, data: Any) -> Any:
+        return data
+
+    def _get_last_metrics(self) -> Optional[Dict[str, Any]]:
+        """Return the last metrics reported by the training function.
+        This function should only be called by LocalController
+        """
+        return self._last_metrics
+
+    def get_all_reported_checkpoints(
+        self,
+        consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
+        timeout_s: Optional[float] = None,
+    ) -> List["ReportedCheckpoint"]:
+        return []
+
+
+_train_fn_utils: Optional[TrainFnUtils] = None
+_train_fn_utils_lock = threading.Lock()
+
+
+def get_train_fn_utils() -> TrainFnUtils:
+    """Return the Ray Train function utilities.
+
+    Returns:
+        The TrainFnUtils instance for the current worker.
+
+    Raises:
+        RuntimeError: If the Ray Train function utilities are not initialized.
+    """
+    global _train_fn_utils
+    with _train_fn_utils_lock:
+        if _train_fn_utils is None:
+            raise RuntimeError("Ray Train function utilities not initialized.")
+        return _train_fn_utils
+
+
+def set_train_fn_utils(train_fn_utils) -> None:
+    global _train_fn_utils
+    with _train_fn_utils_lock:
+        _train_fn_utils = train_fn_utils

python/ray/train/v2/_internal/execution/training_report.py

@@ -0,0 +1,22 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+if TYPE_CHECKING:
+    from ray.train import Checkpoint
+    from ray.train.v2.api.validation_config import ValidationTaskConfig
+
+
+class _TrainingReport:
+    """Checkpoint and metrics reported by user, as well as optional validation configuration."""
+
+    def __init__(
+        self,
+        checkpoint: Optional["Checkpoint"],
+        metrics: Dict[str, Any],
+        validation: Union[bool, "ValidationTaskConfig"],
+    ):
+        self.checkpoint = checkpoint
+        self.metrics = metrics
+        self.validation = validation
+
+    def __repr__(self) -> str:
+        return f"TrainingReport(checkpoint={self.checkpoint}, metrics={self.metrics}, validation={self.validation})"

python/ray/train/v2/_internal/execution/worker_group/__init__.py

@@ -0,0 +1,31 @@
+from .execution_group import ExecutionGroup, ReplicaGroup
+from .placement_group_handle import (
+    DefaultPlacementGroupHandle,
+    PlacementGroupHandle,
+    SlicePlacementGroupHandle,
+)
+from .poll import WorkerGroupPollStatus, WorkerStatus
+from .state import (
+    WorkerGroupContext,
+    WorkerGroupState,
+    WorkerGroupStateBuilder,
+)
+from .worker import ActorMetadata, RayTrainWorker, Worker
+from .worker_group import WorkerGroup
+
+__all__ = [
+    "ActorMetadata",
+    "DefaultPlacementGroupHandle",
+    "ExecutionGroup",
+    "PlacementGroupHandle",
+    "RayTrainWorker",
+    "ReplicaGroup",
+    "SlicePlacementGroupHandle",
+    "Worker",
+    "WorkerGroup",
+    "WorkerGroupContext",
+    "WorkerGroupPollStatus",
+    "WorkerGroupState",
+    "WorkerGroupStateBuilder",
+    "WorkerStatus",
+]

python/ray/train/v2/_internal/execution/worker_group/execution_group.py

@@ -0,0 +1,117 @@
+import abc
+from typing import TYPE_CHECKING, Callable, List, Optional, TypeVar
+
+import ray
+from ray.train._internal.base_worker_group import BaseWorkerGroup
+from ray.train.v2._internal.execution.worker_group.state import _shutdown_workers
+from ray.train.v2._internal.execution.worker_group.worker import Worker
+from ray.types import ObjectRef
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.execution.callback import ReplicaGroupCallback
+    from ray.train.v2._internal.execution.worker_group.state import WorkerGroupContext
+
+T = TypeVar("T")
+
+
+class ExecutionGroup(BaseWorkerGroup):
+    """Base class for groups that can execute functions on workers.
+
+    Provides concrete implementations of the 4 execution methods and __len__
+    based on two abstract primitives: _assert_active() and get_workers().
+    """
+
+    @abc.abstractmethod
+    def _assert_active(self):
+        """Assert that this execution group is active."""
+        pass
+
+    @abc.abstractmethod
+    def get_workers(self) -> List[Worker]:
+        """Return the list of workers in this group."""
+        pass
+
+    def execute_async(self, fn: Callable, *fn_args, **fn_kwargs) -> List[ObjectRef]:
+        self._assert_active()
+        workers = self.get_workers()
+
+        return [worker.execute_async(fn, *fn_args, **fn_kwargs) for worker in workers]
+
+    def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> List[T]:
+        return ray.get(self.execute_async(fn, *fn_args, **fn_kwargs))
+
+    def execute_single_async(
+        self, rank: int, fn: Callable[..., T], *fn_args, **fn_kwargs
+    ) -> ObjectRef:
+        self._assert_active()
+        workers = self.get_workers()
+
+        if rank >= len(workers):
+            raise ValueError(
+                f"The provided {rank=} is " f"not valid for {len(workers)} workers."
+            )
+
+        return workers[rank].execute_async(fn, *fn_args, **fn_kwargs)
+
+    def execute_single(
+        self, rank: int, fn: Callable[..., T], *fn_args, **fn_kwargs
+    ) -> T:
+        return ray.get(self.execute_single_async(rank, fn, *fn_args, **fn_kwargs))
+
+    def __len__(self) -> int:
+        self._assert_active()
+        return len(self.get_workers())
+
+
+class ReplicaGroup(ExecutionGroup):
+    """A group representing a subset of workers from a WorkerGroup.
+
+    Used to pass a replica's workers to backend methods (on_start, etc.)
+    as if they were a standalone worker group.
+    """
+
+    def __init__(
+        self,
+        workers: List[Worker],
+        resources_per_worker: dict,
+        callbacks: Optional[List["ReplicaGroupCallback"]] = None,
+    ):
+        self._workers = workers
+        self._resources_per_worker = resources_per_worker
+        self._callbacks = callbacks or []
+        # An inactive ReplicaGroup still needs to keep track of workers
+        # so we can replace them later.
+        self._active = True
+
+    def _assert_active(self):
+        if not self.is_active():
+            raise ValueError("ReplicaGroup has been shut down.")
+
+    def is_active(self) -> bool:
+        return self._active
+
+    def get_workers(self) -> List[Worker]:
+        return self._workers
+
+    def get_resources_per_worker(self) -> dict:
+        return self._resources_per_worker
+
+    def shutdown(self):
+        """Shutdown all workers in this replica group and clear state."""
+        if self.is_active():
+            for cb in self._callbacks:
+                cb.before_replica_group_shutdown(self)
+
+            _shutdown_workers(self._workers)
+            self._active = False
+
+    def start_training(self, worker_group_context: "WorkerGroupContext"):
+        """Start training on all workers in this replica group."""
+        for cb in self._callbacks:
+            cb.after_replica_group_start(self)
+        ray.get(
+            [
+                worker.actor.run_train_fn.remote(worker_group_context.train_fn_ref)
+                for worker in self._workers
+            ]
+        )

python/ray/train/v2/_internal/execution/worker_group/placement_group_handle.py

@@ -0,0 +1,106 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Union
+
+from ray.types import ObjectRef
+from ray.util.placement_group import PlacementGroup, remove_placement_group
+
+if TYPE_CHECKING:
+    from ray.util.tpu import SlicePlacementGroup
+
+logger = logging.getLogger(__name__)
+
+
+class PlacementGroupHandle(ABC):
+    """Unified interface for placement groups in Ray Train.
+
+    This abstract base class provides a common interface for both standard
+    PlacementGroup and SlicePlacementGroup, allowing WorkerGroup to handle
+    them uniformly without conditional logic.
+    """
+
+    @property
+    @abstractmethod
+    def placement_group(self) -> PlacementGroup:
+        """The underlying PlacementGroup for worker scheduling."""
+        ...
+
+    @abstractmethod
+    def ready(self) -> ObjectRef:
+        """Returns an ObjectRef to check if the placement group is ready.
+
+        Compatible with ray.get() and ray.wait().
+        """
+        ...
+
+    @abstractmethod
+    def wait(self, timeout_seconds: Union[float, int] = 30) -> bool:
+        """Wait for the placement group to be ready within the specified time.
+        Args:
+            timeout_seconds: Timeout in seconds.
+        Returns:
+            True if the placement group is created. False otherwise.
+        """
+        ...
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        """Release all resources associated with this placement group.
+
+        After calling this method, the placement group should no longer be used.
+        """
+        ...
+
+
+class DefaultPlacementGroupHandle(PlacementGroupHandle):
+    """Wrapper for standard PlacementGroup."""
+
+    def __init__(self, pg: PlacementGroup):
+        self._pg = pg
+
+    @property
+    def placement_group(self) -> PlacementGroup:
+        return self._pg
+
+    def ready(self) -> ObjectRef:
+        return self._pg.ready()
+
+    def wait(self, timeout_seconds: Union[float, int] = 30) -> bool:
+        try:
+            return self._pg.wait(timeout_seconds)
+        except Exception:
+            logger.warning(
+                "Placement group wait failed; treating as not ready.",
+                exc_info=True,
+            )
+            return False
+
+    def shutdown(self) -> None:
+        remove_placement_group(self._pg)
+
+
+class SlicePlacementGroupHandle(PlacementGroupHandle):
+    """Wrapper for SlicePlacementGroup that delegates to its underlying PlacementGroup."""
+
+    def __init__(self, spg: "SlicePlacementGroup"):
+        self._spg = spg
+
+    @property
+    def placement_group(self) -> PlacementGroup:
+        return self._spg.placement_group
+
+    def ready(self) -> ObjectRef:
+        return self._spg.placement_group.ready()
+
+    def wait(self, timeout_seconds: Union[float, int] = 30) -> bool:
+        try:
+            return self._spg.placement_group.wait(timeout_seconds)
+        except Exception:
+            logger.warning(
+                "Slice placement group wait failed; treating as not ready.",
+                exc_info=True,
+            )
+            return False
+
+    def shutdown(self) -> None:
+        self._spg.shutdown()

python/ray/train/v2/_internal/execution/worker_group/poll.py

@@ -0,0 +1,152 @@
+import re
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Set
+
+from ray._private.ray_logging import NUMBERS
+from ray.train.v2._internal.exceptions import (
+    UserExceptionWithTraceback,
+    WorkerHealthCheckFailedError,
+)
+from ray.train.v2._internal.execution.preemption import PreemptionInfo
+from ray.train.v2._internal.execution.training_report import _TrainingReport
+from ray.train.v2.api.exceptions import WorkerGroupError
+from ray.types import ObjectRef
+
+ERR_CHAR_LIMIT = 1000
+
+
+def _normalize_error_string(error_str: str) -> str:
+    # Replace numbers with <NUM> based on NUMBERS regex
+    normalized = re.sub(NUMBERS, "<NUM>", error_str)
+    return normalized
+
+
+def _truncate_error_string(error_str: str) -> str:
+    """
+    Truncates error strings to include the first ERR_CHAR_LIMIT // 2
+    characters and the last ERR_CHAR_LIMIT // 2 characters.
+    """
+    if len(error_str) >= ERR_CHAR_LIMIT:
+        return (
+            error_str[: ERR_CHAR_LIMIT // 2]
+            + "...\n... (Output truncated. See individual worker logs for full details) ...\n"
+            + error_str[len(error_str) - ERR_CHAR_LIMIT // 2 :]
+        )
+    return error_str
+
+
+@dataclass
+class WorkerStatus:
+    running: bool
+    error: Optional[Exception] = None
+    training_report: Optional[_TrainingReport] = None
+    return_value: Any = field(default=None)
+    preemption_info: Optional[PreemptionInfo] = None
+
+
+@dataclass(frozen=True)
+class WorkerGroupPollStatus:
+    worker_statuses: Dict[int, WorkerStatus]
+    worker_rank_to_replica_group_rank: Optional[Dict[int, int]] = None
+
+    @property
+    def all_replica_group_indices(self) -> Set[int]:
+        """Return the set of all replica group indices."""
+        if self.worker_rank_to_replica_group_rank is None:
+            return set()
+        return set(self.worker_rank_to_replica_group_rank.values())
+
+    @property
+    def failing_replica_group_indices(self) -> Set[int]:
+        """Return the set of replica group indices that have failing workers."""
+        if self.worker_rank_to_replica_group_rank is None:
+            return set()
+        return {
+            self.worker_rank_to_replica_group_rank[rank]
+            for rank in self.errors
+            if rank in self.worker_rank_to_replica_group_rank
+        }
+
+    @property
+    def errors(self) -> Dict[int, Exception]:
+        errors = {}
+        for world_rank, status in self.worker_statuses.items():
+            if status.error is not None:
+                error = status.error
+                if isinstance(error, UserExceptionWithTraceback):
+                    error = error._base_exc
+                errors[world_rank] = error
+        return errors
+
+    def get_worker_group_error(self) -> WorkerGroupError:
+        return WorkerGroupError(
+            error_message=self.get_error_string(),
+            worker_failures=self.errors,
+        )
+
+    @property
+    def finished(self) -> bool:
+        return self.worker_statuses and all(
+            not status.running for status in self.worker_statuses.values()
+        )
+
+    def get_error_string(self) -> str:
+        """
+        Returns a string representation of worker group errors.
+        Groups similar errors (ignoring numbers) and shows original error examples.
+        """
+        # Group errors by normalized strings (ignoring numbers)
+        normalized_error_to_ranks = defaultdict(list)
+        normalized_error_to_original = {}
+        show_full_error = set()
+
+        for world_rank, status in self.worker_statuses.items():
+            if status.error:
+                error_str = str(status.error)
+                normalized_error = _normalize_error_string(error_str)
+
+                normalized_error_to_ranks[normalized_error].append(str(world_rank))
+
+                # Store the first original error for this normalized group
+                if normalized_error not in normalized_error_to_original:
+                    normalized_error_to_original[normalized_error] = error_str
+
+                # Fully show errors for non-graceful worker failures or running workers
+                if (
+                    isinstance(status.error, WorkerHealthCheckFailedError)
+                    or status.running
+                ):
+                    show_full_error.add(normalized_error)
+
+        errors = []
+        for normalized_error, ranks in normalized_error_to_ranks.items():
+            # Show the original error
+            orig_error = normalized_error_to_original[normalized_error]
+
+            # Convert rank list to comma-separated strings
+            ranks_str = ",".join(ranks)
+
+            if normalized_error in show_full_error:
+                errors.append(f"[Rank {ranks_str} Error Snippet]:\n{orig_error}")
+            else:
+                errors.append(
+                    f"[Rank {ranks_str} Error Snippet]:\n{_truncate_error_string(orig_error)}"
+                )
+
+        error_str = "\n".join(errors)
+
+        return error_str
+
+
+@dataclass(frozen=True)
+class PollTask:
+    """Represents a poll task for a worker.
+
+    Attributes:
+        start_time: The time when the poll task was started.
+        task: The ObjectRef representing the poll task.
+    """
+
+    start_time: float
+    task: ObjectRef

python/ray/train/v2/_internal/execution/worker_group/state.py

@@ -0,0 +1,170 @@
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional
+
+import ray
+from ray.actor import ActorHandle
+from ray.train.v2._internal.execution.checkpoint.sync_actor import SynchronizationActor
+from ray.train.v2._internal.execution.worker_group.worker import Worker
+from ray.train.v2._internal.util import ObjectRefWrapper, time_monotonic
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.execution.worker_group.placement_group_handle import (
+        PlacementGroupHandle,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class WorkerGroupContext:
+    """Context for a worker group.
+
+    This stores the context that is shared when starting a worker group.
+
+    Attributes:
+        run_attempt_id: The ID of the run attempt.
+        train_fn_ref: An object store reference to the training function to execute.
+        num_workers: The number of workers in the worker group.
+        resources_per_worker: The resources per worker.
+        placement_strategy: Strategy for placing workers.
+        label_selector: Optional label selectors to apply per-bundle for workers.
+        num_slices: The number of TPU slices (if using TPU). Defaults to 1.
+    """
+
+    run_attempt_id: str
+    train_fn_ref: ObjectRefWrapper[Callable[[], None]]
+    num_workers: int
+    resources_per_worker: Dict[str, float]
+    placement_strategy: str = "PACK"
+    label_selector: Optional[List[Dict[str, str]]] = None
+    num_slices: int = 1
+
+
+@dataclass(frozen=True)
+class WorkerGroupState:
+    """Ongoing state of an active worker group.
+
+    Attributes:
+        start_time: The time when the worker group was started.
+        workers: The workers in the worker group.
+            These should always be in sorted order by world rank.
+        placement_group_handle: The placement group handle for the worker group.
+        sync_actor: The synchronization actor for the worker group.
+    """
+
+    start_time: float
+    placement_group_handle: "PlacementGroupHandle"
+    workers: List[Worker]
+    sync_actor: ActorHandle
+
+    @property
+    def num_workers(self) -> int:
+        return len(self.workers)
+
+    def replace_workers(
+        self, old_workers: List[Worker], new_workers: List[Worker]
+    ) -> "WorkerGroupState":
+        """Return a new WorkerGroupState with old_workers replaced by new_workers."""
+        current_workers = list(self.workers)
+        for old_w, new_w in zip(old_workers, new_workers):
+            idx = current_workers.index(old_w)
+            current_workers[idx] = new_w
+        return WorkerGroupState(
+            start_time=self.start_time,
+            placement_group_handle=self.placement_group_handle,
+            workers=current_workers,
+            sync_actor=self.sync_actor,
+        )
+
+    def shutdown(self):
+        _shutdown_workers(self.workers)
+        _shutdown_sync_actor(self.sync_actor)
+        self.placement_group_handle.shutdown()
+
+
+class WorkerGroupStateBuilder:
+    """Builder for WorkerGroupState.
+
+    Example usage:
+        ```python
+        builder = WorkerGroupStateBuilder()
+        builder.with_placement_group_handle(placement_group_handle)
+        builder.with_workers(workers)
+        builder.with_sync_actor(sync_actor)
+        state = builder.build()
+
+        builder.shutdown(patience_s=10)
+        ```
+    """
+
+    def __init__(self):
+        self.placement_group_handle = None
+        self.workers = None
+        self.sync_actor = None
+
+    def with_placement_group_handle(
+        self, placement_group_handle: "PlacementGroupHandle"
+    ) -> "WorkerGroupStateBuilder":
+        self.placement_group_handle = placement_group_handle
+        return self
+
+    def with_workers(self, workers: List[Worker]) -> "WorkerGroupStateBuilder":
+        self.workers = workers
+        return self
+
+    def with_sync_actor(
+        self, sync_actor: SynchronizationActor
+    ) -> "WorkerGroupStateBuilder":
+        self.sync_actor = sync_actor
+        return self
+
+    def build(self) -> WorkerGroupState:
+        required_attrs = {
+            "placement_group_handle": self.placement_group_handle,
+            "workers": self.workers,
+            "sync_actor": self.sync_actor,
+        }
+        missing = [name for name, attr in required_attrs.items() if attr is None]
+        if missing:
+            raise ValueError(
+                f"Cannot build incomplete state. Missing: {', '.join(missing)}"
+            )
+        return WorkerGroupState(
+            start_time=time_monotonic(),
+            placement_group_handle=self.placement_group_handle,
+            workers=self.workers,
+            sync_actor=self.sync_actor,
+        )
+
+    def shutdown(self):
+        if self.workers:
+            _shutdown_workers(self.workers)
+            self.workers = None
+
+        if self.sync_actor:
+            _shutdown_sync_actor(self.sync_actor)
+            self.sync_actor = None
+
+        if self.placement_group_handle:
+            self.placement_group_handle.shutdown()
+            self.placement_group_handle = None
+
+
+def _shutdown_workers(workers: List[Worker], patience_s: float = 5):
+    """Shuts down workers after allowing a maximum of patience_s seconds for shutdown hooks to run."""
+    if patience_s < 0:
+        raise ValueError("Invalid patience_s: must be non-negative")
+
+    done_refs = [w.actor.shutdown.remote() for w in workers]
+
+    logger.debug(f"Shutting down {len(workers)} workers.")
+
+    ray.wait(done_refs, num_returns=len(done_refs), timeout=patience_s)
+
+    for worker in workers:
+        ray.kill(worker.actor)
+
+
+def _shutdown_sync_actor(sync_actor: SynchronizationActor):
+    ray.kill(sync_actor)

python/ray/train/v2/_internal/execution/worker_group/thread_runner.py

@@ -0,0 +1,93 @@
+import logging
+import queue
+import threading
+from typing import Callable, Optional, TypeVar
+
+from ray.train.v2._internal.exceptions import UserExceptionWithTraceback
+from ray.train.v2._internal.util import (
+    construct_user_exception_with_traceback,
+    get_callable_name,
+)
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+class ThreadRunner:
+    """Utility to run a user function as a thread and capture its return value
+    or exception.
+    """
+
+    def __init__(self):
+        self._ret: Optional[T] = None
+        self._exc: Optional[UserExceptionWithTraceback] = None
+
+        self._thread: Optional[threading.Thread] = None
+        self._monitor_thread: Optional[threading.Thread] = None
+        self._lock = threading.Lock()
+        self._exc_queue: queue.SimpleQueue[Optional[Exception]] = queue.SimpleQueue()
+
+    def run(self, target: Callable[[], T]) -> None:
+        if self._thread is not None:
+            raise RuntimeError("Thread is already running.")
+
+        def _run_target():
+            try:
+                result = target()
+                with self._lock:
+                    self._ret = result
+                self._exc_queue.put(None)
+            except BaseException as e:
+                # Exclude the first 3 frames from the traceback, which are
+                # the `ThreadRunner._run_target`, `construct_train_func`, and
+                # train_fn_with_final_checkpoint_flush calls.
+                self._exc_queue.put(
+                    construct_user_exception_with_traceback(e, exclude_frames=3)
+                )
+
+            # Join the monitor thread. This ensures that a queued exception
+            # is processed before the target function is considered done.
+            self._monitor_thread.join()
+
+        self._monitor_thread = threading.Thread(
+            target=self._monitor_target,
+            daemon=True,
+            name=f"MonitoringThread({get_callable_name(target)})",
+        )
+        self._monitor_thread.start()
+
+        self._thread = threading.Thread(
+            target=_run_target,
+            daemon=True,
+            name=f"TrainingThread({get_callable_name(target)})",
+        )
+        self._thread.start()
+
+    def _monitor_target(self):
+        """Monitor the exception queue and set the exception if an exception is found.
+
+        This should run as a daemon thread and exit when None is put into the exception queue.
+        """
+        exc: Optional[UserExceptionWithTraceback] = self._exc_queue.get()
+        if exc is None:
+            return
+
+        with self._lock:
+            self._exc = exc
+
+    def is_running(self) -> bool:
+        """Returns whether the target function is still running."""
+        return self._thread is not None and self._thread.is_alive()
+
+    def get_error(self) -> Optional[BaseException]:
+        with self._lock:
+            return self._exc
+
+    def get_return_value(self) -> Optional[T]:
+        with self._lock:
+            return self._ret
+
+    def get_exception_queue(self) -> queue.SimpleQueue:
+        """Returns a queue that nested threads can add exceptions to."""
+        return self._exc_queue

python/ray/train/v2/_internal/execution/worker_group/worker.py

@@ -0,0 +1,319 @@
+import logging
+import os
+import queue
+import socket
+import sys
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, TypeVar, Union
+
+import ray
+import ray._private.ray_constants as ray_constants
+from .thread_runner import ThreadRunner
+from ray.actor import ActorHandle
+from ray.train import Checkpoint
+from ray.train.v2._internal.constants import (
+    DEFAULT_ENABLE_WORKER_LOGGING,
+    ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR,
+)
+from ray.train.v2._internal.execution.callback import (
+    TrainContextCallback,
+    WorkerCallback,
+)
+from ray.train.v2._internal.execution.context import (
+    DistributedContext,
+    ExecutionContext,
+    TrainContext,
+    TrainRunContext,
+    get_train_context,
+    set_train_context,
+)
+from ray.train.v2._internal.execution.preemption import (
+    PreemptionContext,
+    PreemptionInfo,
+)
+from ray.train.v2._internal.execution.storage import StorageContext
+from ray.train.v2._internal.execution.train_fn_utils import (
+    DistributedTrainFnUtils,
+    set_train_fn_utils,
+)
+from ray.train.v2._internal.execution.worker_group.poll import WorkerStatus
+from ray.train.v2._internal.logging.logging import LoggingManager
+from ray.train.v2._internal.logging.patch_print import patch_print_function
+from ray.train.v2._internal.util import ObjectRefWrapper
+from ray.types import ObjectRef
+
+if TYPE_CHECKING:
+    from ray.train.v2._internal.data_integration.interfaces import DatasetShardProvider
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ActorMetadata:
+    hostname: str
+    node_id: str
+    node_ip: str
+    pid: int
+    accelerator_ids: Dict[str, List[Union[int, str]]]
+
+    @property
+    def gpu_ids(self) -> List[Union[int, str]]:
+        return self.accelerator_ids.get("GPU", [])
+
+    @cached_property
+    def _repr(self) -> str:
+        indent = "  "
+        repr_lines = [
+            "ActorMetadata(",
+            f"{indent}hostname={repr(self.hostname)},",
+            f"{indent}node_id={repr(self.node_id)},",
+            f"{indent}node_ip={repr(self.node_ip)},",
+            f"{indent}pid={repr(self.pid)},",
+        ]
+        non_empty_accelerator_ids = {k: v for k, v in self.accelerator_ids.items() if v}
+        if non_empty_accelerator_ids:
+            repr_lines.append(f"{indent}accelerator_ids={non_empty_accelerator_ids},")
+
+        repr_lines.append(")")
+        return "\n".join(repr_lines)
+
+    def __repr__(self) -> str:
+        return self._repr
+
+
+@dataclass
+class Worker:
+    actor: ActorHandle
+    metadata: ActorMetadata
+    resources: Dict[str, float]
+    distributed_context: Optional[DistributedContext] = None
+    log_file_path: Optional[str] = None
+    placement_group_bundle_index: Optional[int] = None
+
+    @cached_property
+    def _repr(self) -> str:
+        indent = "  "
+        metadata_repr = repr(self.metadata).replace("\n", f"\n{indent}")
+        context_repr = repr(self.distributed_context).replace("\n", f"\n{indent}")
+
+        repr_lines = [
+            "Worker(",
+            f"{indent}actor={repr(self.actor)},",
+            f"{indent}metadata={metadata_repr},",
+            f"{indent}distributed_context={context_repr},",
+            f"{indent}log_file_path={repr(self.log_file_path)},",
+            ")",
+        ]
+        return "\n".join(repr_lines)
+
+    def __repr__(self) -> str:
+        return self._repr
+
+    def execute_async(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> ObjectRef:
+        """Execute ``func`` on worker.
+
+        Args:
+            fn: The function to execute on the worker.
+            *fn_args: Positional arguments to forward to ``fn``.
+            **fn_kwargs: Keyword arguments to forward to ``fn``.
+
+        Returns:
+            (ObjectRef) An ObjectRef representing the output of func.
+
+        """
+        return self.actor.execute.options(name=f"execute.{fn.__name__}").remote(
+            fn, *fn_args, **fn_kwargs
+        )
+
+
+class RayTrainWorker:
+    def __init__(self):
+        self._callbacks: List[WorkerCallback] = []
+
+    def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T:
+        return fn(*fn_args, **fn_kwargs)
+
+    def run_train_fn(self, train_fn_ref: ObjectRefWrapper[Callable[[], None]]):
+        """Run the training function in a separate thread.
+
+        This function should return immediately, freeing up the main actor thread
+        to perform other tasks such as polling the status.
+        """
+        try:
+            train_fn = train_fn_ref.get()
+        except Exception as e:
+            logger.error(f"Error deserializing the training function: {e}")
+            raise
+
+        def train_fn_with_final_checkpoint_flush():
+            result = train_fn()
+            get_train_context().checkpoint_upload_threadpool.shutdown()
+
+            if "torch" in sys.modules:
+                from ray.air._internal.torch_utils import contains_tensor
+
+                if contains_tensor(result):
+                    raise ValueError(
+                        "Returning objects containing Torch tensors from the "
+                        "training function is not supported as it will throw an "
+                        "exception on deserialization. You can either convert "
+                        "the tensors to Python objects (ex: `.numpy()`, "
+                        "`.item()`, etc.) or save tensors as part of the "
+                        "checkpoint files instead."
+                    )
+
+            return result
+
+        # Create and start the training thread.
+        logger.debug(
+            f"Rank {get_train_context().get_world_rank()}: Launching training function."
+        )
+        get_train_context().execution_context.training_thread_runner.run(
+            train_fn_with_final_checkpoint_flush
+        )
+
+    def get_metadata(self) -> ActorMetadata:
+        return ActorMetadata(
+            hostname=socket.gethostname(),
+            node_id=ray.get_runtime_context().get_node_id(),
+            node_ip=ray.util.get_node_ip_address(),
+            pid=os.getpid(),
+            accelerator_ids=ray.get_runtime_context().get_accelerator_ids(),
+        )
+
+    def mark_preempt(self, info: PreemptionInfo) -> None:
+        """Store an incoming preemption signal for the UDF to read.
+
+        Called by the PreemptionWatcher on every worker when a preemption
+        affecting the worker group is detected.
+        """
+        train_context = get_train_context()
+        rank = train_context.get_world_rank()
+        train_context.preemption_context.set(info)
+        logger.info(
+            "Rank %d received preemption signal "
+            "(this_worker_preempted=%s, preempted_ranks=%s, deadline_ms=%s).",
+            rank,
+            rank in info.preempted_ranks,
+            info.preempted_ranks,
+            info.deadline_ms,
+        )
+
+    def poll_status(self) -> WorkerStatus:
+        train_context = get_train_context()
+        execution_context = train_context.execution_context
+
+        # TODO: We can implement two phase commit here.
+        # Only mark the task done when the result has been processed by the controller.
+        try:
+            training_report = execution_context.result_queue.get_nowait()
+            execution_context.result_queue.task_done()
+        except queue.Empty:
+            training_report = None
+
+        error = execution_context.training_thread_runner.get_error()
+
+        # TODO: The running state should not be conflated with queue flushing.
+        # Running should only be true if the user code is still running.
+        # This relies on `worker_group_status.finished` returning False
+        # until all training results have been flushed.
+        running = execution_context.training_thread_runner.is_running() or bool(
+            training_report
+        )
+
+        return_value = (
+            execution_context.training_thread_runner.get_return_value()
+            if not running
+            else None
+        )
+
+        return WorkerStatus(
+            running=running,
+            error=error,
+            training_report=training_report,
+            return_value=return_value,
+            preemption_info=train_context.preemption_context.get(),
+        )
+
+    def clear_result_queue(self) -> bool:
+        """Drain the result queue, discarding any pending training reports.
+
+        Returns:
+            True if the queue had at least one result, False if it was empty.
+        """
+        execution_context = get_train_context().execution_context
+        had_result = False
+        while True:
+            try:
+                execution_context.result_queue.get_nowait()
+                execution_context.result_queue.task_done()
+                had_result = True
+            except queue.Empty:
+                break
+        return had_result
+
+    def shutdown(self):
+        """Shutdown the worker.
+
+        This method is not doing the real shutdown, but it is used by the worker
+        group to signal the worker to stop running the training function.
+        Any shutdown worker callbacks can hook on this method to implement the
+        corresponding shutdown logic. Note that the shutdown logic needs to be
+        thread-safe if it is running in a separate thread.
+        """
+        for callback in self._callbacks:
+            callback.before_worker_shutdown()
+
+    def init_train_context(
+        self,
+        train_run_context: TrainRunContext,
+        distributed_context: DistributedContext,
+        synchronization_actor: ActorHandle,
+        storage_context: StorageContext,
+        worker_callbacks: List[Union[WorkerCallback, TrainContextCallback]],
+        controller_actor: ActorHandle,
+        dataset_shard_provider: Optional["DatasetShardProvider"] = None,
+        checkpoint: Optional[Checkpoint] = None,
+        has_validation_fn: Optional[bool] = None,
+        current_report_index: int = 0,
+    ):
+        self._callbacks = [c for c in worker_callbacks if isinstance(c, WorkerCallback)]
+        context_callbacks_to_propagate = [
+            c for c in worker_callbacks if isinstance(c, TrainContextCallback)
+        ]
+        context = TrainContext(
+            train_run_context=train_run_context,
+            distributed_context=distributed_context,
+            execution_context=ExecutionContext(
+                synchronization_actor=synchronization_actor,
+                # Make the queue size 1 to avoid building up too
+                # many unprocessed results.
+                result_queue=queue.Queue(maxsize=1),
+                training_thread_runner=ThreadRunner(),
+                train_context_callbacks=context_callbacks_to_propagate,
+            ),
+            storage_context=storage_context,
+            preemption_context=PreemptionContext(),
+            controller_actor=controller_actor,
+            checkpoint=checkpoint,
+            dataset_shard_provider=dataset_shard_provider,
+            has_validation_fn=has_validation_fn,
+            current_report_index=current_report_index,
+        )
+        # Configure the train and root logger for the worker processes.
+        if ray_constants.env_bool(
+            ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR, DEFAULT_ENABLE_WORKER_LOGGING
+        ):
+            LoggingManager.configure_worker_logger(context)
+        patch_print_function()
+        # Set the train context global variable for the worker.
+        set_train_context(context)
+
+        # user facing train fn utils
+        set_train_fn_utils(DistributedTrainFnUtils())
+
+        for callback in self._callbacks:
+            callback.after_init_train_context()

python/ray/train/v2/_internal/logging/__init__.py

@@ -0,0 +1,3 @@
+from .logging import LoggingManager
+
+__all__ = ["LoggingManager"]

python/ray/train/v2/_internal/logging/logging.py

@@ -0,0 +1,340 @@
+import logging.config
+import os
+from enum import Enum
+from typing import Optional, Union
+
+import ray
+from ray._common.filters import CoreContextFilter
+from ray._common.formatters import JSONFormatter
+from ray._private.log import PlainRayHandler
+from ray.train.v2._internal.execution.context import TrainContext, TrainRunContext
+from ray.train.v2._internal.util import get_module_name
+
+
+class TrainContextFilter(logging.Filter):
+    """Add Ray Train metadata to the log records.
+
+    This filter is applied to Ray Train controller and worker processes.
+    """
+
+    # Log keys for Ray Train controller and worker processes.
+    class LogKey(str, Enum):
+        RUN_NAME = "run_name"
+        COMPONENT = "component"
+        WORLD_RANK = "world_rank"
+        LOCAL_RANK = "local_rank"
+        NODE_RANK = "node_rank"
+
+    # Ray Train Component by process types
+    class TrainComponent(str, Enum):
+        CONTROLLER = "controller"
+        WORKER = "worker"
+
+    def __init__(self, context: Union[TrainRunContext, TrainContext]):
+        self._is_worker: bool = isinstance(context, TrainContext)
+        if self._is_worker:
+            self._run_name: str = context.train_run_context.get_run_config().name
+            self._world_rank: int = context.get_world_rank()
+            self._local_rank: int = context.get_local_rank()
+            self._node_rank: int = context.get_node_rank()
+            self._component: str = TrainContextFilter.TrainComponent.WORKER
+        else:
+            self._run_name: str = context.get_run_config().name
+            self._component: str = TrainContextFilter.TrainComponent.CONTROLLER
+
+    def controller_filter(self, record):
+        # Add the run_id and component to Ray Train controller processes.
+        setattr(record, TrainContextFilter.LogKey.RUN_NAME, self._run_name)
+        setattr(record, TrainContextFilter.LogKey.COMPONENT, self._component)
+        return True
+
+    def worker_filter(self, record):
+        # Add the run_id and component to Ray Train worker processes.
+        setattr(record, TrainContextFilter.LogKey.RUN_NAME, self._run_name)
+        setattr(record, TrainContextFilter.LogKey.COMPONENT, self._component)
+        # Add all the rank related information to the log record for worker processes.
+        setattr(record, TrainContextFilter.LogKey.WORLD_RANK, self._world_rank)
+        setattr(record, TrainContextFilter.LogKey.LOCAL_RANK, self._local_rank)
+        setattr(record, TrainContextFilter.LogKey.NODE_RANK, self._node_rank)
+        return True
+
+    def filter(self, record):
+        if self._is_worker:
+            return self.worker_filter(record)
+        else:
+            return self.controller_filter(record)
+
+
+class TrainLogLevelFilter(logging.Filter):
+    """Filter that applies log level filtering only to ray.train log records."""
+
+    def __init__(self, log_level: str = "INFO"):
+        super().__init__()
+        self._log_level = getattr(logging, log_level)
+
+    def filter(self, record):
+        if record.name == "ray.train" or record.name.startswith("ray.train."):
+            return record.levelno >= self._log_level
+        return True
+
+
+class SessionFileHandler(logging.Handler):
+    """A handler that writes to a log file in the Ray session directory.
+
+    The Ray session directory isn't available until Ray is initialized, so any logs
+    emitted before Ray is initialized will be lost.
+    This handler will not create the file handler until you emit a log record.
+
+    Args:
+        filename: The name of the log file. The file is created in the 'logs/train'
+            directory of the Ray session directory.
+    """
+
+    # TODO (hpguo): This handler class is shared by both Ray Train and ray data. We
+    # should move this to ray core and make it available to both libraries.
+
+    def __init__(self, filename: str):
+        super().__init__()
+        self._filename = filename
+        self._handler = None
+        self._formatter = None
+        self._path = None
+
+    def emit(self, record):
+        if self._handler is None:
+            self._try_create_handler()
+        if self._handler is not None:
+            self._handler.emit(record)
+
+    def setFormatter(self, fmt: logging.Formatter) -> None:
+        if self._handler is not None:
+            self._handler.setFormatter(fmt)
+        self._formatter = fmt
+
+    def get_log_file_path(self) -> Optional[str]:
+        if self._handler is None:
+            self._try_create_handler()
+        return self._path
+
+    def _try_create_handler(self):
+        assert self._handler is None
+
+        # Get the Ray Train log directory. If not in a Ray session, return.
+        # This handler will only be created within a Ray session.
+        log_directory = LoggingManager.get_log_directory()
+        if log_directory is None:
+            return
+
+        os.makedirs(log_directory, exist_ok=True)
+
+        # Create the log file.
+        self._path = os.path.join(log_directory, self._filename)
+        self._handler = logging.FileHandler(self._path)
+        if self._formatter is not None:
+            self._handler.setFormatter(self._formatter)
+
+
+class LoggingManager:
+    """
+    A utility class for managing the logging configuration of Ray Train.
+    """
+
+    @staticmethod
+    def _get_base_logger_config_dict(
+        context: Union[TrainRunContext, TrainContext],
+    ) -> dict:
+        """Return the base logging configuration dictionary."""
+        log_level = LoggingManager._resolve_log_level(context)
+        # Using Ray worker ID as the file identifier where logs are written to.
+        file_identifier = ray.get_runtime_context().get_worker_id()
+        # Return the base logging configuration as a Python dictionary.
+        return {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "ray_json": {"class": get_module_name(JSONFormatter)},
+            },
+            "filters": {
+                "core_context_filter": {"()": CoreContextFilter},
+                "train_context_filter": {"()": TrainContextFilter, "context": context},
+                "train_log_level_filter": {
+                    "()": TrainLogLevelFilter,
+                    "log_level": log_level,
+                },
+            },
+            "handlers": {
+                "console": {
+                    "class": get_module_name(PlainRayHandler),
+                    "filters": ["train_log_level_filter"],
+                },
+                "file_train_sys_controller": {
+                    "class": get_module_name(SessionFileHandler),
+                    "formatter": "ray_json",
+                    "filename": f"ray-train-sys-controller-{file_identifier}.log",
+                    "filters": ["core_context_filter", "train_context_filter"],
+                },
+                "file_train_app_controller": {
+                    "class": get_module_name(SessionFileHandler),
+                    "formatter": "ray_json",
+                    "filename": f"ray-train-app-controller-{file_identifier}.log",
+                    "filters": [
+                        "core_context_filter",
+                        "train_context_filter",
+                        "train_log_level_filter",
+                    ],
+                },
+                "file_train_sys_worker": {
+                    "class": get_module_name(SessionFileHandler),
+                    "formatter": "ray_json",
+                    "filename": f"ray-train-sys-worker-{file_identifier}.log",
+                    "filters": ["core_context_filter", "train_context_filter"],
+                },
+                "file_train_app_worker": {
+                    "class": get_module_name(SessionFileHandler),
+                    "formatter": "ray_json",
+                    "filename": f"ray-train-app-worker-{file_identifier}.log",
+                    "filters": [
+                        "core_context_filter",
+                        "train_context_filter",
+                        "train_log_level_filter",
+                    ],
+                },
+            },
+            "loggers": {},
+        }
+
+    @staticmethod
+    def _resolve_log_level(
+        context: Union[TrainRunContext, TrainContext],
+    ) -> str:
+        """Returns the log level from RunConfig's LoggingConfig."""
+        if isinstance(context, TrainContext):
+            run_config = context.train_run_context.get_run_config()
+        else:
+            run_config = context.get_run_config()
+
+        return run_config.logging_config.log_level
+
+    @staticmethod
+    def _get_controller_logger_config_dict(context: TrainRunContext) -> dict:
+        """Return the controller logger configuration dictionary.
+
+        On the controller process, only the `ray.train` logger is configured.
+        It is broadly set to level DEBUG, with downstream processing by log handlers.
+        This logger emits logs to the following three locations:
+            - `file_train_sys_controller`: Ray Train system logs.
+            - `file_train_app_controller`: Ray Train application logs.
+            - `console`: Logs to the console.
+        """
+
+        config_dict = LoggingManager._get_base_logger_config_dict(context)
+        config_dict["loggers"]["ray.train"] = {
+            "level": "DEBUG",
+            "handlers": [
+                "file_train_sys_controller",
+                "file_train_app_controller",
+                "console",
+            ],
+            "propagate": False,
+        }
+        return config_dict
+
+    @staticmethod
+    def _get_worker_logger_config_dict(context: TrainContext) -> dict:
+        """Return the worker loggers configuration dictionary.
+
+        On the worker process, there are two loggers being configured:
+
+        First, the `ray.train` logger is configured and emits logs to the
+        following three locations:
+            - `file_train_sys_worker`: Ray Train system logs.
+            - `file_train_app_worker`: Ray Train application logs.
+            - `console`: Logs to the console.
+        It is broadly set to level DEBUG, with downstream processing by log handlers.
+
+        Second, the root logger is configured and emits logs to the following
+        two locations:
+            - `console`: Logs to the console.
+            - `file_train_app_worker`: Ray Train application logs.
+        The root logger will not emit Ray Train system logs and thus not writing to
+        `file_train_sys_worker` file handler.
+        """
+
+        config_dict = LoggingManager._get_base_logger_config_dict(context)
+        config_dict["loggers"]["ray.train"] = {
+            "level": "DEBUG",
+            "handlers": ["file_train_sys_worker", "file_train_app_worker", "console"],
+            "propagate": False,
+        }
+        config_dict["root"] = {
+            "level": "INFO",
+            "handlers": ["file_train_app_worker", "console"],
+        }
+        return config_dict
+
+    @staticmethod
+    def configure_controller_logger(context: TrainRunContext) -> None:
+        """
+        Configure the logger on the controller process, which is the `ray.train` logger.
+        """
+        config = LoggingManager._get_controller_logger_config_dict(context)
+        logging.config.dictConfig(config)
+        # TODO: Return the controller log file path.
+
+    @staticmethod
+    def configure_worker_logger(context: TrainContext) -> None:
+        """
+        Configure the loggers on the worker process, which contains the
+        `ray.train` logger and the root logger.
+        """
+        config = LoggingManager._get_worker_logger_config_dict(context)
+        logging.config.dictConfig(config)
+        # TODO: Return the worker log file path.
+
+    @staticmethod
+    def get_log_directory() -> Optional[str]:
+        """Return the directory where Ray Train writes log files.
+
+        If not in a Ray session, return None.
+
+        This path looks like: "/tmp/ray/session_xxx/logs/train/"
+        """
+        global_node = ray._private.worker._global_node
+
+        if global_node is None:
+            return None
+
+        root_dir = global_node.get_session_dir_path()
+        return os.path.join(root_dir, "logs", "train")
+
+
+def get_train_application_controller_log_path() -> Optional[str]:
+    """
+    Return the path to the file train application controller log file.
+    """
+    # TODO: This is a temporary solution. We should return the log file path in
+    # the `configure_controller_logger` function.
+    logger = logging.getLogger("ray.train")
+    for handler in logger.handlers:
+        if (
+            isinstance(handler, SessionFileHandler)
+            and "ray-train-app-controller" in handler._filename
+        ):
+            return handler.get_log_file_path()
+    return None
+
+
+def get_train_application_worker_log_path() -> Optional[str]:
+    """
+    Return the path to the file train application worker log file.
+    """
+    # TODO: This is a temporary solution. We should return the log file path in
+    # the `configure_worker_logger` function.
+    logger = logging.getLogger("ray.train")
+    for handler in logger.handlers:
+        if (
+            isinstance(handler, SessionFileHandler)
+            and "ray-train-app-worker" in handler._filename
+        ):
+            return handler.get_log_file_path()
+    return None

python/ray/train/v2/_internal/logging/patch_print.py

@@ -0,0 +1,76 @@
+import builtins
+import contextlib
+import logging
+import sys
+from typing import Callable
+
+from ray._private.ray_constants import env_bool
+from ray.train.v2._internal.constants import (
+    DEFAULT_ENABLE_PRINT_PATCH,
+    ENABLE_PRINT_PATCH_ENV_VAR,
+)
+
+# Save the original print function
+_original_print = builtins.print
+
+
+@contextlib.contextmanager
+def print_context_manager(print_fn: Callable):
+    """Context manager to set the builtin print function as print_fn."""
+    current_print = builtins.print
+    builtins.print = print_fn
+    yield
+    builtins.print = current_print
+
+
+def redirected_print(*objects, sep=" ", end="\n", file=None, flush=False):
+    """Implement python's print function to redirect logs to Train's logger.
+
+    If the file is set to anything other than stdout, stderr, or None, call the
+    builtin print. Else, construct the message and redirect to Train's logger.
+
+    This makes sure that print to customized file in user defined function will not
+    be overwritten by the redirected print function.
+
+    See https://docs.python.org/3/library/functions.html#print
+    """
+    # TODO (hpguo): This handler class is shared by both ray train and ray serve. We
+    # should move this to ray core and make it available to both libraries.
+
+    if file not in [sys.stdout, sys.stderr, None]:
+        _original_print(*objects, sep=sep, end=end, file=file, flush=flush)
+        return
+
+    # If sys.stdout/stderr has been redirected (e.g. contextlib.redirect_stdout(),
+    # or wrapping by libraries like wandb / MLflow / colorama / IPython), tee to
+    # the original print so the redirect target also receives the output. The
+    # logger still gets the message below, so structured logs aren't silently
+    # dropped when a third-party library wraps the stream.
+    if (file in (sys.stdout, None) and sys.stdout is not sys.__stdout__) or (
+        file is sys.stderr and sys.stderr is not sys.__stderr__
+    ):
+        _original_print(*objects, sep=sep, end=end, file=file, flush=flush)
+
+    root_logger = logging.getLogger()
+    message = sep.join(map(str, objects))
+    # Use the original `print` method for the scope of the logger call, in order to
+    # avoid infinite recursion errors if any exceptions get raised (since exception
+    # handling involves another `print(..., file=sys.stderr)`.
+    # Note that an exception being raised here is not expected (e.g. it would be a
+    # bug in our own logging code), so this is just to keep the error logs sane
+    # during development.
+    with print_context_manager(_original_print):
+        # We want this log to be associated with the line of code where user calls
+        # `print`, which is stacklevel 2.
+        # Frame [stacklevel]:
+        # User's call to print [2] -> `redirected_print` [1] -> root_logger.log [0]
+        root_logger.log(logging.INFO, message, stacklevel=2)
+
+
+def patch_print_function() -> None:
+    """
+    Patch the print function to redirect logs to Train's logger.
+    Only patch the print function if the environment variable is set to "1"
+    """
+    if env_bool(ENABLE_PRINT_PATCH_ENV_VAR, DEFAULT_ENABLE_PRINT_PATCH):
+        builtins.print = redirected_print

python/ray/train/v2/_internal/metrics/base.py

@@ -0,0 +1,166 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Dict, Generic, Optional, Tuple, TypeVar
+
+from ray.util.metrics import Gauge
+
+RUN_NAME_TAG_KEY = "ray_train_run_name"
+RUN_ID_TAG_KEY = "ray_train_run_id"
+
+T = TypeVar("T")
+E = TypeVar("E", bound=Enum)
+
+
+class Metric(ABC):
+    def __init__(
+        self,
+        name: str,
+        default: T,
+        description: str,
+        base_tags: Dict[str, str],
+    ):
+        """
+        Initialize a new metric.
+
+        Args:
+            name: The name of the metric.
+            default: The default value of the metric.
+            description: The description of the metric.
+            base_tags: The base tags for the metric.
+        """
+        self._default = default
+        self._base_tags = base_tags
+        self._gauge = Gauge(
+            name,
+            description=description,
+            tag_keys=self._get_tag_keys(),
+        )
+
+    @abstractmethod
+    def record(self, value: T):
+        """Update the metric value.
+
+        Args:
+            value: The value to update the metric with.
+        """
+        pass
+
+    @abstractmethod
+    def get_value(self) -> T:
+        """Get the value of the metric.
+
+        Returns:
+            The value of the metric. If the metric has not been recorded,
+            the default value is returned.
+        """
+        pass
+
+    @abstractmethod
+    def reset(self):
+        """Reset values and clean up resources."""
+        pass
+
+    def _get_tag_keys(self) -> Tuple[str, ...]:
+        return tuple(self._base_tags.keys())
+
+
+class TimeMetric(Metric):
+    """A metric for tracking elapsed time."""
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        base_tags: Dict[str, str],
+    ):
+        self._current_value = 0.0
+        super().__init__(
+            name=name,
+            default=0.0,
+            description=description,
+            base_tags=base_tags,
+        )
+
+    def record(self, value: float):
+        """Update the time metric value by accumulating the time.
+
+        Args:
+            value: The time value to increment the metric by.
+        """
+        self._current_value += value
+        self._gauge.set(self._current_value, self._base_tags)
+
+    def get_value(self) -> float:
+        return self._current_value
+
+    def reset(self):
+        self._current_value = self._default
+        self._gauge.set(self._default, self._base_tags)
+
+
+class EnumMetric(Metric, Generic[E]):
+    """A metric for tracking enum values."""
+
+    DEFAULT_VALUE = 0
+    RECORDED_VALUE = 1
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        base_tags: Dict[str, str],
+        enum_tag_key: str,
+    ):
+        self._enum_tag_key = enum_tag_key
+        self._current_value: Optional[E] = None
+        super().__init__(
+            name=name,
+            default=self.DEFAULT_VALUE,
+            description=description,
+            base_tags=base_tags,
+        )
+
+    def record(self, enum_value: E) -> None:
+        """Record a specific enum value.
+
+        The metric will be reset to 0 for the previous value and set to 1 for the new value.
+
+        Args:
+            enum_value: The enum value to record for.
+        """
+        if enum_value == self._current_value:
+            return
+
+        if self._current_value is not None:
+            previous_tags = self._get_tags(self._current_value)
+            self._gauge.set(self._default, previous_tags)
+
+        current_tags = self._get_tags(enum_value)
+        self._gauge.set(self.RECORDED_VALUE, current_tags)
+
+        self._current_value = enum_value
+
+    def get_value(self, enum_value: E) -> int:
+        """Get the value for a specific enum value.
+
+        Args:
+            enum_value: The enum value to get the value for
+
+        Returns:
+            The value for the enum value
+        """
+        return int(enum_value == self._current_value)
+
+    def reset(self):
+        if self._current_value is not None:
+            tags = self._get_tags(self._current_value)
+            self._gauge.set(self._default, tags)
+        self._current_value = None
+
+    def _get_tag_keys(self) -> Tuple[str, ...]:
+        return tuple(self._base_tags.keys()) + (self._enum_tag_key,)
+
+    def _get_tags(self, enum_value: E) -> Dict[str, str]:
+        tags = self._base_tags.copy()
+        tags[self._enum_tag_key] = enum_value.name
+        return tags

python/ray/train/v2/_internal/metrics/controller.py

@@ -0,0 +1,66 @@
+from typing import Dict, Union
+
+from ray.train.v2._internal.execution.controller.state import TrainControllerStateType
+from ray.train.v2._internal.metrics.base import (
+    RUN_ID_TAG_KEY,
+    RUN_NAME_TAG_KEY,
+    EnumMetric,
+    TimeMetric,
+)
+
+
+class ControllerMetrics:
+    """Factory for creating controller-specific metrics.
+
+    This class defines all metrics used to track the state and performance of the
+    training controller. Each metric is defined with its name, type, default value,
+    description, and required tags.
+    """
+
+    # ===== Metric Names =====
+    CONTROLLER_STATE = "train_controller_state"
+    WORKER_GROUP_START_TOTAL_TIME_S = "train_worker_group_start_total_time_s"
+    WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S = "train_worker_group_shutdown_total_time_s"
+
+    # ===== Tag Keys =====
+    CONTROLLER_STATE_TAG_KEY = "ray_train_controller_state"
+
+    @classmethod
+    def _create_time_metric(
+        cls, name: str, description: str, base_tags: Dict[str, str]
+    ) -> TimeMetric:
+        return TimeMetric(
+            name=name,
+            description=description,
+            base_tags=base_tags,
+        )
+
+    @classmethod
+    def _create_controller_state_metric(
+        cls, base_tags: Dict[str, str]
+    ) -> EnumMetric[TrainControllerStateType]:
+        return EnumMetric[TrainControllerStateType](
+            name=cls.CONTROLLER_STATE,
+            description="Current state of the Ray Train controller",
+            base_tags=base_tags,
+            enum_tag_key=cls.CONTROLLER_STATE_TAG_KEY,
+        )
+
+    @classmethod
+    def get_controller_metrics(
+        cls, run_name: str, run_id: str
+    ) -> Dict[str, Union[TimeMetric, EnumMetric[TrainControllerStateType]]]:
+        base_tags = {RUN_NAME_TAG_KEY: run_name, RUN_ID_TAG_KEY: run_id}
+        return {
+            cls.WORKER_GROUP_START_TOTAL_TIME_S: cls._create_time_metric(
+                cls.WORKER_GROUP_START_TOTAL_TIME_S,
+                "Total time taken to start the worker group",
+                base_tags,
+            ),
+            cls.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S: cls._create_time_metric(
+                cls.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S,
+                "Total time taken to shutdown the worker group",
+                base_tags,
+            ),
+            cls.CONTROLLER_STATE: cls._create_controller_state_metric(base_tags),
+        }

python/ray/train/v2/_internal/metrics/worker.py

@@ -0,0 +1,64 @@
+from typing import Dict
+
+from ray.train.v2._internal.metrics.base import (
+    RUN_ID_TAG_KEY,
+    RUN_NAME_TAG_KEY,
+    TimeMetric,
+)
+
+WORKER_WORLD_RANK_TAG_KEY = "ray_train_worker_world_rank"
+WORKER_ACTOR_ID_TAG_KEY = "ray_train_worker_actor_id"
+
+
+class WorkerMetrics:
+    """Factory for creating worker-specific metrics.
+
+    This class defines all metrics used to track the state and performance of the
+    training workers. Each metric is defined with its name, type, default value,
+    description, and required tags.
+    """
+
+    # ===== Metric Names =====
+    REPORT_TOTAL_BLOCKED_TIME_S = "train_report_total_blocked_time_s"
+    CHECKPOINT_SYNC_TOTAL_TIME_S = "train_checkpoint_sync_total_time_s"
+    CHECKPOINT_TRANSFER_TOTAL_TIME_S = "train_checkpoint_transfer_total_time_s"
+
+    @classmethod
+    def _create_time_metric(
+        cls, name: str, description: str, base_tags: Dict[str, str]
+    ) -> TimeMetric:
+        """Create a time-based metric."""
+        return TimeMetric(
+            name=name,
+            description=description,
+            base_tags=base_tags,
+        )
+
+    @classmethod
+    def get_worker_metrics(
+        cls, run_name: str, run_id: str, world_rank: int, worker_actor_id: str
+    ) -> Dict[str, TimeMetric]:
+        """Get all worker metrics."""
+        base_tags = {
+            RUN_NAME_TAG_KEY: run_name,
+            RUN_ID_TAG_KEY: run_id,
+            WORKER_WORLD_RANK_TAG_KEY: str(world_rank),
+            WORKER_ACTOR_ID_TAG_KEY: worker_actor_id,
+        }
+        return {
+            cls.REPORT_TOTAL_BLOCKED_TIME_S: cls._create_time_metric(
+                cls.REPORT_TOTAL_BLOCKED_TIME_S,
+                "Cumulative time in seconds to report a checkpoint to the storage.",
+                base_tags,
+            ),
+            cls.CHECKPOINT_SYNC_TOTAL_TIME_S: cls._create_time_metric(
+                cls.CHECKPOINT_SYNC_TOTAL_TIME_S,
+                "Cumulative time in seconds spent synchronizing the checkpoint directory name across all ranks.",
+                base_tags,
+            ),
+            cls.CHECKPOINT_TRANSFER_TOTAL_TIME_S: cls._create_time_metric(
+                cls.CHECKPOINT_TRANSFER_TOTAL_TIME_S,
+                "Cumulative time in seconds spent transferring checkpoint files to storage.",
+                base_tags,
+            ),
+        }

python/ray/train/v2/_internal/migration_utils.py

@@ -0,0 +1,70 @@
+from ray.train.constants import V2_MIGRATION_GUIDE_MESSAGE
+
+FAIL_FAST_DEPRECATION_MESSAGE = (
+    "`ray.train.FailureConfig(fail_fast)` is deprecated since it is "
+    "only relevant in the context of multiple trials running in Ray Tune. "
+    "This parameter is still available in `ray.tune.FailureConfig` "
+    "for passing into a `ray.tune.Tuner`. "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+TRAINER_RESOURCES_DEPRECATION_MESSAGE = (
+    "`ray.train.ScalingConfig(trainer_resources)` is deprecated. "
+    "This parameter was an advanced configuration that specified "
+    "resources for the Ray Train driver actor, which doesn't "
+    "need to reserve logical resources because it doesn't perform "
+    "any heavy computation. "
+    "Only the `resources_per_worker` parameter should be used "
+    "to specify resources for the training workers. "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+VERBOSE_DEPRECATION_MESSAGE = (
+    "`ray.train.RunConfig(verbose)` is deprecated. "
+    "This parameter controls Ray Tune logging verbosity, "
+    "and is only relevant when using Ray Tune. "
+    "This parameter is still available in `ray.tune.RunConfig` "
+    "for passing into a `ray.tune.Tuner`. "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+LOG_TO_FILE_DEPRECATION_MESSAGE = (
+    "`ray.train.RunConfig(log_to_file)` is deprecated. "
+    "The Ray Train driver actor and the training worker actors "
+    "already log stdout/stderr as part of Ray's logging system. "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+STOP_DEPRECATION_MESSAGE = (
+    "`ray.train.RunConfig(stop)` is deprecated. "
+    "This parameter is only relevant when using Ray Tune "
+    "and is still available in `ray.tune.RunConfig` "
+    "for passing into a `ray.tune.Tuner`. "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+CALLBACKS_DEPRECATION_MESSAGE = (
+    "`ray.train.RunConfig(callbacks: List[ray.tune.Callback])` is deprecated. "
+    "Ray Train no longer accepts Ray Tune callbacks, since the Ray Train "
+    "execution backend is being separated from Ray Tune. "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+PROGRESS_REPORTER_DEPRECATION_MESSAGE = (
+    "`ray.train.RunConfig(progress_reporter)` is deprecated. "
+    "This parameter controls the Ray Tune console output reporter, "
+    "and is only relevant when using Ray Tune. "
+    "This parameter is still available in `ray.tune.RunConfig` "
+    "for passing into a `ray.tune.Tuner`. "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)
+
+SYNC_CONFIG_DEPRECATION_MESSAGE = (
+    "`ray.train.RunConfig(sync_config)` is deprecated. "
+    "This configuration controls advanced syncing behavior, "
+    "which is either not supported or not relevant in the reworked Ray Train. "
+    "This parameter is still available in `ray.tune.RunConfig` "
+    "for passing into a `ray.tune.Tuner`. "
+    "The `SyncConfig` class has been moved to `ray.tune.SyncConfig`. "
+    f"{V2_MIGRATION_GUIDE_MESSAGE}"
+)

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

@@ -0,0 +1,313 @@
+import logging
+
+from google.protobuf.struct_pb2 import Struct
+
+from ray.core.generated.export_train_state_pb2 import (
+    ExportTrainRunAttemptEventData as ProtoTrainRunAttempt,
+    ExportTrainRunEventData as ProtoTrainRun,
+)
+from ray.dashboard.modules.metrics.dashboards.common import Panel
+from ray.dashboard.modules.metrics.dashboards.train_dashboard_panels import (
+    TRAIN_RUN_PANELS,
+    TRAIN_WORKER_PANELS,
+)
+from ray.train.v2._internal.state.schema import (
+    ActorStatus,
+    BackendConfig,
+    DataConfig,
+    ExecutionOptions,
+    RunAttemptStatus,
+    RunConfig,
+    RunSettings,
+    RunStatus,
+    ScalingConfig,
+    TrainRun,
+    TrainRunAttempt,
+    TrainWorker,
+)
+from ray.train.v2._internal.util import TrainingFramework
+
+# Increment each time the exported Train schema changes (proto, pydantic, or
+# exported json) so downstream consumers can distinguish schema versions.
+TRAIN_SCHEMA_VERSION = 4
+RAY_TRAIN_VERSION = 2
+
+# Status mapping dictionaries
+_ACTOR_STATUS_MAP = {
+    ActorStatus.ALIVE: ProtoTrainRunAttempt.ActorStatus.ALIVE,
+    ActorStatus.DEAD: ProtoTrainRunAttempt.ActorStatus.DEAD,
+}
+
+_RUN_ATTEMPT_STATUS_MAP = {
+    RunAttemptStatus.PENDING: ProtoTrainRunAttempt.RunAttemptStatus.PENDING,
+    RunAttemptStatus.RUNNING: ProtoTrainRunAttempt.RunAttemptStatus.RUNNING,
+    RunAttemptStatus.FINISHED: ProtoTrainRunAttempt.RunAttemptStatus.FINISHED,
+    RunAttemptStatus.ERRORED: ProtoTrainRunAttempt.RunAttemptStatus.ERRORED,
+    RunAttemptStatus.ABORTED: ProtoTrainRunAttempt.RunAttemptStatus.ABORTED,
+}
+
+_RUN_STATUS_MAP = {
+    RunStatus.INITIALIZING: ProtoTrainRun.RunStatus.INITIALIZING,
+    RunStatus.SCHEDULING: ProtoTrainRun.RunStatus.SCHEDULING,
+    RunStatus.RUNNING: ProtoTrainRun.RunStatus.RUNNING,
+    RunStatus.RESTARTING: ProtoTrainRun.RunStatus.RESTARTING,
+    RunStatus.RESIZING: ProtoTrainRun.RunStatus.RESIZING,
+    RunStatus.FINISHED: ProtoTrainRun.RunStatus.FINISHED,
+    RunStatus.ERRORED: ProtoTrainRun.RunStatus.ERRORED,
+    RunStatus.ABORTED: ProtoTrainRun.RunStatus.ABORTED,
+}
+
+_TRAINING_FRAMEWORK_MAP = {
+    None: ProtoTrainRun.BackendConfig.TrainingFramework.TRAINING_FRAMEWORK_UNSPECIFIED,
+    TrainingFramework.TORCH: ProtoTrainRun.BackendConfig.TrainingFramework.TORCH,
+    TrainingFramework.JAX: ProtoTrainRun.BackendConfig.TrainingFramework.JAX,
+    TrainingFramework.TENSORFLOW: ProtoTrainRun.BackendConfig.TrainingFramework.TENSORFLOW,
+    TrainingFramework.XGBOOST: ProtoTrainRun.BackendConfig.TrainingFramework.XGBOOST,
+    TrainingFramework.LIGHTGBM: ProtoTrainRun.BackendConfig.TrainingFramework.LIGHTGBM,
+}
+
+logger = logging.getLogger(__name__)
+
+
+def _dict_to_struct(d: dict) -> Struct:
+    """Returns a protobuf Struct from a dictionary."""
+    s = Struct()
+    s.update(d)
+    return s
+
+
+def _to_proto_resources(resources: dict) -> ProtoTrainRunAttempt.TrainResources:
+    """Convert resources dictionary to protobuf TrainResources."""
+    return ProtoTrainRunAttempt.TrainResources(resources=resources)
+
+
+def _to_proto_worker(worker: TrainWorker) -> ProtoTrainRunAttempt.TrainWorker:
+    """Convert TrainWorker to protobuf format."""
+    status = None
+    if worker.status is not None:
+        status = _ACTOR_STATUS_MAP[worker.status]
+
+    return 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=status,
+        resources=_to_proto_resources(worker.resources.resources),
+        log_file_path=worker.log_file_path,
+    )
+
+
+# Main conversion functions
+def train_run_attempt_to_proto(attempt: TrainRunAttempt) -> ProtoTrainRunAttempt:
+    """Convert TrainRunAttempt to protobuf format."""
+    proto_attempt = ProtoTrainRunAttempt(
+        schema_version=TRAIN_SCHEMA_VERSION,
+        ray_train_version=RAY_TRAIN_VERSION,
+        run_id=attempt.run_id,
+        attempt_id=attempt.attempt_id,
+        status=_RUN_ATTEMPT_STATUS_MAP[attempt.status],
+        status_detail=attempt.status_detail,
+        start_time_ns=attempt.start_time_ns,
+        end_time_ns=attempt.end_time_ns,
+        resources=[_to_proto_resources(r.resources) for r in attempt.resources],
+        workers=[_to_proto_worker(w) for w in attempt.workers],
+    )
+
+    return proto_attempt
+
+
+def _to_proto_dashboard_panel(panel: Panel) -> ProtoTrainRun.DashboardPanelMetadata:
+    """Convert Dashboard Panel to protobuf format."""
+    proto_panel = ProtoTrainRun.DashboardPanelMetadata(
+        id=str(panel.id),
+        title=panel.title,
+    )
+
+    return proto_panel
+
+
+def to_proto_backend_config(
+    backend_config: BackendConfig,
+) -> ProtoTrainRun.BackendConfig:
+    """Convert BackendConfig to protobuf format."""
+    proto_backend_config = ProtoTrainRun.BackendConfig(
+        framework=_TRAINING_FRAMEWORK_MAP[backend_config.framework],
+    )
+
+    proto_backend_config.config.CopyFrom(_dict_to_struct(backend_config.config))
+
+    return proto_backend_config
+
+
+def to_proto_scaling_config(
+    scaling_config: ScalingConfig,
+) -> ProtoTrainRun.ScalingConfig:
+    """Convert ScalingConfig to protobuf format."""
+    proto_scaling_config = ProtoTrainRun.ScalingConfig(
+        use_gpu=scaling_config.use_gpu,
+        placement_strategy=scaling_config.placement_strategy,
+        use_tpu=scaling_config.use_tpu,
+    )
+
+    if isinstance(scaling_config.num_workers, tuple):
+        proto_scaling_config.num_workers_range.CopyFrom(
+            ProtoTrainRun.ScalingConfig.IntRange(
+                min=scaling_config.num_workers[0],
+                max=scaling_config.num_workers[1],
+            )
+        )
+    else:
+        proto_scaling_config.num_workers_fixed = scaling_config.num_workers
+
+    if scaling_config.resources_per_worker is not None:
+        proto_scaling_config.resources_per_worker.values.update(
+            scaling_config.resources_per_worker
+        )
+
+    if scaling_config.accelerator_type is not None:
+        proto_scaling_config.accelerator_type = scaling_config.accelerator_type
+    if scaling_config.topology is not None:
+        proto_scaling_config.topology = scaling_config.topology
+
+    if scaling_config.bundle_label_selector is not None:
+        selectors = scaling_config.bundle_label_selector
+        if isinstance(selectors, dict):
+            proto_scaling_config.label_selector_single.values.update(selectors)
+        else:
+            proto_scaling_config.label_selector_list.values.extend(
+                [ProtoTrainRun.ScalingConfig.StringMap(values=s) for s in selectors]
+            )
+
+    return proto_scaling_config
+
+
+def _to_proto_execution_options(
+    execution_options: ExecutionOptions,
+) -> ProtoTrainRun.ExecutionOptions:
+    """Convert a single ExecutionOptions schema model to protobuf."""
+    return ProtoTrainRun.ExecutionOptions(
+        resource_limits=_dict_to_struct(execution_options.resource_limits),
+        exclude_resources=_dict_to_struct(execution_options.exclude_resources),
+        preserve_order=execution_options.preserve_order,
+        actor_locality_enabled=execution_options.actor_locality_enabled,
+        verbose_progress=execution_options.verbose_progress,
+    )
+
+
+def to_proto_data_config(data_config: DataConfig) -> ProtoTrainRun.DataConfig:
+    """Convert DataConfig to protobuf format."""
+    data_execution_options = data_config.data_execution_options
+    proto_data_config = ProtoTrainRun.DataConfig(
+        enable_shard_locality=data_config.enable_shard_locality,
+        data_execution_options=ProtoTrainRun.DataExecutionOptions(
+            default=_to_proto_execution_options(data_execution_options.default),
+            per_dataset_execution_options={
+                name: _to_proto_execution_options(opts)
+                for name, opts in data_execution_options.per_dataset_execution_options.items()
+            },
+        ),
+    )
+
+    if data_config.datasets_to_split == "all":
+        proto_data_config.all.SetInParent()
+    else:
+        proto_data_config.datasets.values.extend(data_config.datasets_to_split)
+
+    return proto_data_config
+
+
+def _to_proto_failure_config(run_config: RunConfig) -> ProtoTrainRun.FailureConfig:
+    """Convert RunConfig.failure_config to protobuf format."""
+    return ProtoTrainRun.FailureConfig(
+        max_failures=run_config.failure_config.max_failures,
+        controller_failure_limit=run_config.failure_config.controller_failure_limit,
+    )
+
+
+def _to_proto_checkpoint_config(
+    run_config: RunConfig,
+) -> ProtoTrainRun.CheckpointConfig:
+    """Convert RunConfig.checkpoint_config to protobuf format."""
+    checkpoint_score_order = ProtoTrainRun.CheckpointConfig.CheckpointScoreOrder.Value(
+        run_config.checkpoint_config.checkpoint_score_order.upper()
+    )
+
+    proto_checkpoint_config = ProtoTrainRun.CheckpointConfig(
+        checkpoint_score_order=checkpoint_score_order
+    )
+    if run_config.checkpoint_config.num_to_keep is not None:
+        proto_checkpoint_config.num_to_keep = run_config.checkpoint_config.num_to_keep
+    if run_config.checkpoint_config.checkpoint_score_attribute is not None:
+        proto_checkpoint_config.checkpoint_score_attribute = (
+            run_config.checkpoint_config.checkpoint_score_attribute
+        )
+    return proto_checkpoint_config
+
+
+def to_proto_run_config(run_config: RunConfig) -> ProtoTrainRun.RunConfig:
+    """Convert RunConfig to protobuf format."""
+    proto_run_config = ProtoTrainRun.RunConfig(
+        name=run_config.name,
+        failure_config=_to_proto_failure_config(run_config),
+        worker_runtime_env=_dict_to_struct(run_config.worker_runtime_env),
+        checkpoint_config=_to_proto_checkpoint_config(run_config),
+        storage_path=run_config.storage_path,
+    )
+
+    if run_config.storage_filesystem is not None:
+        proto_run_config.storage_filesystem = run_config.storage_filesystem
+
+    return proto_run_config
+
+
+def _to_proto_run_settings(run_settings: RunSettings) -> ProtoTrainRun.RunSettings:
+    """Convert RunSettings to protobuf format."""
+
+    proto_run_settings = ProtoTrainRun.RunSettings(
+        backend_config=to_proto_backend_config(run_settings.backend_config),
+        scaling_config=to_proto_scaling_config(run_settings.scaling_config),
+        datasets=run_settings.datasets,
+        data_config=to_proto_data_config(run_settings.data_config),
+        run_config=to_proto_run_config(run_settings.run_config),
+    )
+
+    if run_settings.train_loop_config is not None:
+        proto_run_settings.train_loop_config.CopyFrom(
+            _dict_to_struct(run_settings.train_loop_config)
+        )
+
+    return proto_run_settings
+
+
+def train_run_to_proto(run: TrainRun) -> ProtoTrainRun:
+    """Convert TrainRun to protobuf format."""
+
+    proto_train_run_panels = [_to_proto_dashboard_panel(p) for p in TRAIN_RUN_PANELS]
+    proto_train_worker_panels = [
+        _to_proto_dashboard_panel(p) for p in TRAIN_WORKER_PANELS
+    ]
+
+    proto_train_run = ProtoTrainRun(
+        schema_version=TRAIN_SCHEMA_VERSION,
+        ray_train_version=RAY_TRAIN_VERSION,
+        id=run.id,
+        name=run.name,
+        job_id=bytes.fromhex(run.job_id),
+        controller_actor_id=bytes.fromhex(run.controller_actor_id),
+        status=_RUN_STATUS_MAP[run.status],
+        status_detail=run.status_detail,
+        start_time_ns=run.start_time_ns,
+        end_time_ns=run.end_time_ns,
+        controller_log_file_path=run.controller_log_file_path,
+        train_run_panels=proto_train_run_panels,
+        train_worker_panels=proto_train_worker_panels,
+        framework_versions=run.framework_versions,
+        run_settings=_to_proto_run_settings(run.run_settings),
+    )
+
+    return proto_train_run

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

@@ -0,0 +1,532 @@
+import math
+from collections.abc import Mapping
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+from pydantic import field_validator
+
+from ray._common.pydantic_compat import BaseModel, Field
+from ray.dashboard.modules.job.pydantic_models import JobDetails
+from ray.train.v2._internal.util import TrainingFramework
+from ray.util.annotations import DeveloperAPI
+
+MAX_ERROR_STACK_TRACE_LENGTH = 50000
+
+
+def _to_json_serializable_value(value: Any, *, max_depth: int = 3) -> Any:
+    """Recursively coerce a value into a human-readable, JSON serializable representation.
+
+    If ``value`` is a list or dict, this function walks through it and replaces non-JSON
+    serializable fields (e.g. custom objects, modules, tensors, callables, etc.) with a
+    human-readable string representation.
+
+    Args:
+        value: Any Python value. Primitives pass through; collections recurse;
+            other types are stringified.
+        max_depth: Truncates dicts nested beyond ``max_depth`` to ``"..."``.
+            Lists do not consume depth.
+
+    Returns:
+        The JSON serializable representation of the value.
+    """
+    if max_depth <= 0:
+        raise ValueError("max_depth must be greater than 0")
+
+    def _safe_str(v):
+        try:
+            return str(v)
+        except Exception:
+            return type(v).__name__
+
+    def _walk(value, depth):
+        if value is None or isinstance(value, (bool, int, str)):
+            return value
+        if isinstance(value, float):
+            return str(value) if not math.isfinite(value) else value
+        if isinstance(value, Mapping):
+            if depth <= 0:
+                return "..."
+            try:
+                items = list(value.items())
+            except Exception:
+                # Custom Mapping subclass with a broken `.items()`.
+                return type(value).__name__
+            return {_safe_str(k): _walk(v, depth - 1) for k, v in items}
+
+        # Tuples, sets, and frozensets all become lists in JSON.
+        if isinstance(value, (list, tuple, set, frozenset)):
+            return [_walk(v, depth) for v in value]
+
+        cls = type(value)
+        # Use class name if no custom string representation is defined.
+        if cls.__str__ is object.__str__ and cls.__repr__ is object.__repr__:
+            return cls.__name__
+        return _safe_str(value)
+
+    return _walk(value, max_depth)
+
+
+@DeveloperAPI
+class RunStatus(str, Enum):
+    """Enumeration of the possible statuses for a Train run."""
+
+    # ====== Active States ======
+    # The Train run is currently in the process of initializing.
+    INITIALIZING = "INITIALIZING"
+    # The Train run is waiting to be scheduled.
+    SCHEDULING = "SCHEDULING"
+    # The Train run is currently in progress.
+    RUNNING = "RUNNING"
+    # The Train run is recovering from a failure or restart.
+    RESTARTING = "RESTARTING"
+    # The Train run is resizing.
+    RESIZING = "RESIZING"
+
+    # ===== Terminal States ======
+    # The Train run completed successfully.
+    FINISHED = "FINISHED"
+    # The Train run failed due to an error in the training workers.
+    ERRORED = "ERRORED"
+    # The Train run was terminated due to system or controller errors.
+    ABORTED = "ABORTED"
+
+    def is_terminal(self) -> bool:
+        return self in [RunStatus.FINISHED, RunStatus.ERRORED, RunStatus.ABORTED]
+
+
+@DeveloperAPI
+class RunAttemptStatus(str, Enum):
+    """Enumeration of the possible statuses for a Train run attempt."""
+
+    # ====== Active States ======
+    # The run attempt is waiting to be scheduled.
+    PENDING = "PENDING"
+    # The run attempt is currently in progress.
+    RUNNING = "RUNNING"
+
+    # ===== Terminal States =====
+    # The run attempt completed successfully.
+    FINISHED = "FINISHED"
+    # The run attempt failed due to an error in the training workers.
+    ERRORED = "ERRORED"
+    # The run attempt was terminated due to system or controller errors.
+    ABORTED = "ABORTED"
+
+    def is_terminal(self) -> bool:
+        return self in [
+            RunAttemptStatus.FINISHED,
+            RunAttemptStatus.ERRORED,
+            RunAttemptStatus.ABORTED,
+        ]
+
+
+@DeveloperAPI
+class ActorStatus(str, Enum):
+    """Enumeration of the statuses for a Train worker actor."""
+
+    # The actor is currently active.
+    ALIVE = "ALIVE"
+    # The actor is no longer active.
+    DEAD = "DEAD"
+
+
+@DeveloperAPI
+class TrainResources(BaseModel):
+    """Resources allocated for a Train worker or run."""
+
+    resources: Dict[str, float] = Field(
+        description="A dictionary specifying the types and amounts of resources "
+        "allocated (e.g., CPU, GPU)."
+    )
+
+
+@DeveloperAPI
+class TrainWorker(BaseModel):
+    """Metadata about a Ray Train worker."""
+
+    world_rank: int = Field(
+        description="The global rank of the worker in the training cluster."
+    )
+    local_rank: int = Field(description="The local rank of the worker on its node.")
+    node_rank: int = Field(description="The rank of the worker's node in the cluster.")
+    actor_id: str = Field(description="The unique ID of the worker's actor.")
+    node_id: str = Field(
+        description="The unique ID of the node where the worker is running."
+    )
+    node_ip: str = Field(
+        description="The IP address of the node where the worker is running."
+    )
+    pid: int = Field(description="The process ID of the worker.")
+    gpu_ids: List[int] = Field(description="A list of GPU IDs allocated to the worker.")
+    status: Optional[ActorStatus] = Field(
+        None, description="The current status of the worker actor."
+    )
+    resources: TrainResources = Field(
+        description="The resources allocated to this Train worker."
+    )
+    log_file_path: Optional[str] = Field(
+        None, description="The path to the log file for the Train worker."
+    )
+
+
+@DeveloperAPI
+class MemoryInfo(BaseModel):
+    """Memory usage information for a process."""
+
+    rss: int = Field(description="The resident set size (RSS) memory usage in bytes.")
+    vms: int = Field(description="The virtual memory size (VMS) usage in bytes.")
+    pfaults: Optional[int] = Field(None, description="The number of page faults.")
+    pageins: Optional[int] = Field(None, description="The number of page-ins.")
+
+
+@DeveloperAPI
+class ProcessStats(BaseModel):
+    """CPU and memory statistics for a process."""
+
+    cpuPercent: float = Field(description="The percentage of CPU usage.")
+    mem: Optional[List[int]] = Field(
+        None,
+        description="Memory statistics, including total memory, free memory, "
+        "and memory usage ratio.",
+    )
+    memoryInfo: MemoryInfo = Field(description="Detailed memory usage information.")
+
+
+class ProcessGPUUsage(BaseModel):
+    """GPU usage statistics for a process."""
+
+    pid: int = Field(description="The process ID.")
+    gpuMemoryUsage: int = Field(description="The GPU memory usage in bytes.")
+
+
+@DeveloperAPI
+class GPUStats(BaseModel):
+    """Statistics for a GPU."""
+
+    uuid: str = Field(description="The unique identifier of the GPU.")
+    index: int = Field(description="The index of the GPU.")
+    name: str = Field(description="The name of the GPU.")
+    utilizationGpu: Optional[float] = Field(
+        None, description="The percentage utilization of the GPU."
+    )
+    memoryUsed: float = Field(description="The amount of GPU memory used in bytes.")
+    memoryTotal: float = Field(description="The total amount of GPU memory in bytes.")
+    processInfo: ProcessGPUUsage = Field(
+        description="GPU usage statistics for the associated process."
+    )
+
+
+@DeveloperAPI
+class DecoratedTrainWorker(TrainWorker):
+    """Detailed metadata for a Ray Train worker, including process and GPU stats."""
+
+    processStats: Optional[ProcessStats] = Field(
+        None, description="CPU and memory statistics for the worker process."
+    )
+    gpus: List[GPUStats] = Field(
+        default_factory=list,
+        description="A list of GPUs used by the worker process,"
+        " with detailed statistics.",
+    )
+
+
+@DeveloperAPI
+class TrainRunAttempt(BaseModel):
+    """Metadata for an individual attempt to execute a Train run."""
+
+    run_id: str = Field(description="Unique identifier for the parent Train run.")
+    attempt_id: str = Field(
+        description="Unique identifier for this specific Train run attempt."
+    )
+    status: RunAttemptStatus = Field(
+        description="The current execution status of the Train run attempt."
+    )
+    status_detail: Optional[str] = Field(
+        None,
+        description="Additional details about the status,"
+        " including error messages if applicable.",
+    )
+    start_time_ns: int = Field(
+        description="The UNIX timestamp (in nanoseconds)"
+        " when the Train run attempt started."
+    )
+    end_time_ns: Optional[int] = Field(
+        None,
+        description="The UNIX timestamp (in nanoseconds)"
+        " when the Train run attempt ended. "
+        "If null, the attempt is still ongoing.",
+    )
+    resources: List[TrainResources] = Field(
+        description="The resources (e.g., CPU, GPU) allocated to the Train run attempt."
+    )
+    workers: List[TrainWorker] = Field(
+        description="List of Train workers participating in this attempt, "
+        "sorted by global ranks."
+    )
+
+
+@DeveloperAPI
+class DecoratedTrainRunAttempt(TrainRunAttempt):
+    """Detailed metadata for a Train run attempt, including decorated worker data."""
+
+    workers: List[DecoratedTrainWorker] = Field(
+        description="A list of Train workers with detailed statistics, "
+        "sorted by global ranks."
+    )
+
+
+@DeveloperAPI
+class ExecutionOptions(BaseModel):
+    """ExecutionOptions for a single Ray Data ingest pipeline."""
+
+    resource_limits: Dict[str, Any] = Field(
+        description="The resource limits applied to the Ray Data execution plan."
+    )
+    exclude_resources: Dict[str, Any] = Field(
+        description="The resources excluded from the Ray Data execution plan "
+        "(e.g. resources reserved by Ray Train workers)."
+    )
+
+    @field_validator("resource_limits", "exclude_resources", mode="before")
+    @classmethod
+    def _sanitize_dict(cls, v):
+        return _to_json_serializable_value(v)
+
+    preserve_order: bool = Field(
+        description="Whether to preserve the order of outputs across operators."
+    )
+    actor_locality_enabled: bool = Field(
+        description="Whether actor-based locality optimizations are enabled."
+    )
+    verbose_progress: bool = Field(
+        description="Whether verbose progress reporting is enabled."
+    )
+
+
+@DeveloperAPI
+class DataExecutionOptions(BaseModel):
+    """ExecutionOptions for a Ray Train run, split into defaults and per-dataset overrides."""
+
+    default: ExecutionOptions = Field(
+        description="Execution options applied to any dataset without a per-dataset override."
+    )
+    per_dataset_execution_options: Dict[str, ExecutionOptions] = Field(
+        default_factory=dict,
+        description="Per-dataset execution option overrides, keyed by dataset name.",
+    )
+
+
+@DeveloperAPI
+class DataConfig(BaseModel):
+    """Configuration for dataset splitting and execution options within Ray Train."""
+
+    datasets_to_split: Union[Literal["all"], List[str]] = Field(
+        description="Which datasets to split; either 'all' or a list of dataset names."
+    )
+    execution_options: Optional[Dict] = Field(
+        default=None,
+        deprecated="DEPRECATED: Use data_execution_options instead.",
+    )
+    data_execution_options: DataExecutionOptions = Field(
+        description="Data execution options"
+    )
+    enable_shard_locality: bool = Field(
+        description="Whether to enable shard locality optimization."
+    )
+
+
+@DeveloperAPI
+class ScalingConfig(BaseModel):
+    """Scaling config for a Train run."""
+
+    num_workers: Union[int, Tuple[int, int]] = Field(
+        description="The number of workers for the Train run."
+    )
+    use_gpu: bool = Field(description="Whether to use GPUs for the Train run.")
+    resources_per_worker: Optional[Dict[str, float]] = Field(
+        None, description="The resources per worker for a Train run."
+    )
+    placement_strategy: str = Field(
+        description="The placement strategy for the Train run."
+    )
+    accelerator_type: Optional[str] = Field(
+        None, description="The accelerator type for the Train run."
+    )
+    use_tpu: bool = Field(description="Whether to use TPUs for the Train run.")
+    topology: Optional[str] = Field(None, description="The topology for the Train run.")
+    bundle_label_selector: Optional[
+        Union[Dict[str, str], List[Dict[str, str]]]
+    ] = Field(None, description="The bundle label selector for the Train run.")
+
+
+@DeveloperAPI
+class FailureConfig(BaseModel):
+    """Failure config for a Train run."""
+
+    max_failures: int = Field(
+        description="The maximum number of failures for a Train run."
+    )
+    controller_failure_limit: int = Field(
+        description="The maximum number of controller failures to tolerate."
+    )
+
+
+@DeveloperAPI
+class CheckpointConfig(BaseModel):
+    """Checkpoint config for a Train run."""
+
+    num_to_keep: Optional[int] = Field(
+        None,
+        description="The number of most recent checkpoints to keep. Older checkpoints may be deleted.",
+    )
+    checkpoint_score_attribute: Optional[str] = Field(
+        None,
+        description="Attribute used to score and rank checkpoints; can be a metric key or attribute.",
+    )
+    checkpoint_score_order: Literal["max", "min"] = Field(
+        description="Order to rank checkpoint scores, 'max' for higher-is-better, 'min' for lower-is-better.",
+    )
+
+
+@DeveloperAPI
+class RunConfig(BaseModel):
+    """Run configuration parameters for a Train run, encompassing failure,
+    runtime environment, checkpoint settings, and storage path."""
+
+    name: str = Field(description="The name of the Train run.")
+    failure_config: FailureConfig = Field(
+        description="The failure config for a Train run."
+    )
+    worker_runtime_env: Dict[str, Any] = Field(
+        description="The worker runtime env for a Train run."
+    )
+
+    @field_validator("worker_runtime_env", mode="before")
+    @classmethod
+    def _sanitize_worker_runtime_env(cls, v):
+        return _to_json_serializable_value(v)
+
+    checkpoint_config: CheckpointConfig = Field(
+        description="The checkpoint config for a Train run."
+    )
+    storage_path: str = Field(description="The storage path for a Train run.")
+    storage_filesystem: Optional[str] = Field(
+        None, description="The storage filesystem for a Train run."
+    )
+
+    @field_validator("storage_filesystem", mode="before")
+    @classmethod
+    def _sanitize_storage_filesystem(cls, v):
+        return _to_json_serializable_value(v)
+
+
+@DeveloperAPI
+class BackendConfig(BaseModel):
+    """Backend config for a Train run."""
+
+    framework: Optional[TrainingFramework] = Field(
+        None, description="The training framework for this backend config."
+    )
+    config: Dict[str, Any] = Field(
+        description="Training framework-specific configuration fields."
+    )
+
+    @field_validator("config", mode="before")
+    @classmethod
+    def _sanitize_config(cls, v):
+        return _to_json_serializable_value(v)
+
+
+@DeveloperAPI
+class RunSettings(BaseModel):
+    """Settings for a Train run, primarily consisting of configs set before a train run starts.
+
+    This includes the train loop config, backend config, scaling config, dataset configs,
+    and runtime configuration.
+    """
+
+    train_loop_config: Optional[Dict] = Field(
+        None, description="The user defined train loop config for a Train run."
+    )
+
+    @field_validator("train_loop_config", mode="before")
+    @classmethod
+    def _sanitize_train_loop_config(cls, v):
+        return _to_json_serializable_value(v)
+
+    backend_config: BackendConfig = Field(
+        description="The backend config for a Train run. Can vary with the framework (e.g. TorchConfig)"
+    )
+    scaling_config: ScalingConfig = Field(
+        description="The scaling config for this Train run."
+    )
+    datasets: List[str] = Field(
+        description="A list of dataset names for a Train run.",
+    )
+    data_config: DataConfig = Field(
+        description="The data config for a Train run.",
+    )
+    run_config: RunConfig = Field(
+        description="Run configuration for this Train run, including failure, runtime environment, checkpoint settings, and storage path."
+    )
+
+
+@DeveloperAPI
+class TrainRun(BaseModel):
+    """Metadata for a Ray Train run, including its details and status."""
+
+    id: str = Field(description="Unique identifier for the Train run.")
+    name: str = Field(description="Human-readable name assigned to the Train run.")
+    job_id: str = Field(description="The Ray Job ID associated with this Train run.")
+    controller_actor_id: str = Field(
+        description="Unique ID of the actor managing the Train run."
+    )
+    status: RunStatus = Field(
+        description="The current execution status of the Train run."
+    )
+    status_detail: Optional[str] = Field(
+        None,
+        description="Additional details about the current status, "
+        "including error messages if applicable.",
+    )
+    start_time_ns: int = Field(
+        description="The UNIX timestamp (in nanoseconds) when the Train run started."
+    )
+    end_time_ns: Optional[int] = Field(
+        None,
+        description="The UNIX timestamp (in nanoseconds) when the Train run ended. "
+        "If null, the run is still in progress.",
+    )
+    controller_log_file_path: Optional[str] = Field(
+        None, description="The path to the log file for the Train run controller."
+    )
+    framework_versions: Dict[str, str] = Field(
+        description="The relevant framework versions for this Train run,"
+        "including the Ray version and training framework version."
+    )
+    run_settings: RunSettings = Field(
+        description="The run settings for this Train run, including train loop config, "
+        "backend config, scaling config, dataset details, and runtime configuration."
+    )
+
+
+@DeveloperAPI
+class DecoratedTrainRun(TrainRun):
+    """Detailed metadata for a Ray Train run, including attempts and job details."""
+
+    attempts: List[DecoratedTrainRunAttempt] = Field(
+        description="A list of attempts made to execute the Train run."
+    )
+    job_details: Optional[JobDetails] = Field(
+        None,
+        description="Detailed information about the job that initiated this Train run.",
+    )
+
+
+@DeveloperAPI
+class TrainRunsResponse(BaseModel):
+    """Response containing a list of decorated Train runs."""
+
+    train_runs: List[DecoratedTrainRun] = Field(
+        description="A list of Train runs with detailed metadata."
+    )

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

@@ -0,0 +1,301 @@
+import copy
+import logging
+import os
+import threading
+import time
+from collections import OrderedDict, defaultdict
+from typing import Dict, Optional
+
+import ray
+from ray._private import ray_constants
+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.v2._internal.constants import (
+    CONTROLLERS_TO_POLL_PER_ITERATION,
+    DEFAULT_ENABLE_STATE_ACTOR_RECONCILIATION,
+    DEFAULT_STATE_ACTOR_RECONCILIATION_INTERVAL_S,
+    ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
+    GET_ACTOR_TIMEOUT_S,
+    STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR,
+)
+from ray.train.v2._internal.state.schema import (
+    TrainRun,
+    TrainRunAttempt,
+)
+from ray.train.v2._internal.state.util import (
+    is_actor_alive,
+    update_train_run_aborted,
+    update_train_run_attempt_aborted,
+)
+from ray.train.v2._internal.util import time_monotonic
+
+logger = logging.getLogger(__name__)
+
+
+class TrainStateActor:
+    def __init__(
+        self,
+        # TODO: group into single config if we need to do similar polling elsewhere
+        enable_state_actor_reconciliation: bool = False,
+        reconciliation_interval_s: float = 30,
+        get_actor_timeout_s: int = GET_ACTOR_TIMEOUT_S,
+        controllers_to_poll_per_iteration: int = CONTROLLERS_TO_POLL_PER_ITERATION,
+    ):
+        # NOTE: All runs and attempts are stored in memory.
+        # This may be a memory issue for large runs.
+        # TODO: consider cleaning up runs over time.
+        self._runs: Dict[str, TrainRun] = OrderedDict()
+        # {run_id: {attempt_id: TrainRunAttempt}}
+        self._run_attempts: Dict[str, OrderedDict[str, TrainRunAttempt]] = defaultdict(
+            OrderedDict
+        )
+        (
+            self._export_logger,
+            self._is_train_run_export_api_enabled,
+            self._is_train_run_attempt_export_api_enabled,
+        ) = self._init_export_logger()
+
+        # TODO: consider row level locking if loop takes too long.
+        self._runs_lock = threading.RLock()
+        self._run_attempts_lock = threading.RLock()
+
+        # Set env vars related to reconciling train run/attempt state.
+        if enable_state_actor_reconciliation:
+            self._reconciliation_interval_s = reconciliation_interval_s
+            self._controllers_to_poll_per_iteration = controllers_to_poll_per_iteration
+            self._get_actor_timeout_s = get_actor_timeout_s
+            self._start_run_state_reconciliation_thread()
+
+    def _abort_live_runs_with_dead_controllers(
+        self, last_poll_run_id: Optional[str]
+    ) -> str:
+        aborted_run_ids = []
+        with self._runs_lock:
+            runs = list(self._runs.values())
+
+            # Start iterating from poll index.
+            starting_poll_index = 0
+            if last_poll_run_id is not None:
+                for poll_index, run in enumerate(runs):
+                    if run.id == last_poll_run_id:
+                        starting_poll_index = (poll_index + 1) % len(runs)
+                        break
+
+            # Abort runs.
+            num_polled_runs = 0
+            poll_index = starting_poll_index
+            while (
+                poll_index < starting_poll_index + len(runs)
+                and num_polled_runs < self._controllers_to_poll_per_iteration
+            ):
+                run = runs[poll_index % len(runs)]
+                poll_index += 1
+                last_poll_run_id = run.id
+                if run.status.is_terminal():
+                    continue
+                try:
+                    if not is_actor_alive(
+                        run.controller_actor_id, self._get_actor_timeout_s
+                    ):
+                        update_train_run_aborted(run, False)
+                        self.create_or_update_train_run(run)
+                        aborted_run_ids.append(run.id)
+                except ray.util.state.exception.RayStateApiException:
+                    logger.exception(
+                        "State API unavailable when checking if actor is alive. "
+                        "Will check again on next poll."
+                    )
+                num_polled_runs += 1
+
+        # Abort run attempts.
+        with self._run_attempts_lock:
+            for run_id in aborted_run_ids:
+                latest_run_attempt = self._get_latest_run_attempt(run_id)
+                if latest_run_attempt and not latest_run_attempt.status.is_terminal():
+                    update_train_run_attempt_aborted(latest_run_attempt, False)
+                    self.create_or_update_train_run_attempt(latest_run_attempt)
+
+        return last_poll_run_id
+
+    def _start_run_state_reconciliation_thread(self) -> None:
+        def _reconciliation_loop():
+            last_poll_run_id = None
+            latest_poll_time = float("-inf")
+            while True:
+                # Wait for the poll interval to elapse.
+                time_since_last_poll = time_monotonic() - latest_poll_time
+                if time_since_last_poll < self._reconciliation_interval_s:
+                    remaining_time = (
+                        self._reconciliation_interval_s - time_since_last_poll
+                    )
+                    time.sleep(remaining_time)
+
+                last_poll_run_id = self._abort_live_runs_with_dead_controllers(
+                    last_poll_run_id
+                )
+                latest_poll_time = time_monotonic()
+
+        threading.Thread(target=_reconciliation_loop, daemon=True).start()
+
+    def _get_latest_run_attempt(self, run_id: str) -> Optional[TrainRunAttempt]:
+        with self._run_attempts_lock:
+            # NOTE: run_attempts is OrderedDict from attempt_id to TrainRunAttempt.
+            run_attempts = self._run_attempts.get(run_id, {})
+            if not run_attempts:
+                return None
+            return next(reversed(run_attempts.values()))
+
+    def create_or_update_train_run(self, run: TrainRun) -> None:
+        with self._runs_lock:
+            self._runs[run.id] = run
+            run_copy = copy.deepcopy(run)
+        self._maybe_export_train_run(run_copy)
+
+    def create_or_update_train_run_attempt(self, run_attempt: TrainRunAttempt) -> None:
+        with self._run_attempts_lock:
+            self._run_attempts[run_attempt.run_id][run_attempt.attempt_id] = run_attempt
+            run_attempt_copy = copy.deepcopy(run_attempt)
+        self._maybe_export_train_run_attempt(run_attempt_copy)
+
+    def get_train_runs(self) -> Dict[str, TrainRun]:
+        with self._runs_lock:
+            return self._runs
+
+    def get_train_run_attempts(self) -> Dict[str, Dict[str, TrainRunAttempt]]:
+        with self._run_attempts_lock:
+            return self._run_attempts
+
+    # ============================
+    # 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: TrainRun) -> None:
+        if not self._is_train_run_export_api_enabled:
+            return
+
+        from ray.train.v2._internal.state.export import train_run_to_proto
+
+        run_proto = train_run_to_proto(run)
+        self._export_logger.send_event(run_proto)
+
+    def _maybe_export_train_run_attempt(self, run_attempt: TrainRunAttempt) -> None:
+        if not self._is_train_run_attempt_export_api_enabled:
+            return
+
+        from ray.train.v2._internal.state.export import train_run_attempt_to_proto
+
+        run_attempt_proto = train_run_attempt_to_proto(run_attempt)
+        self._export_logger.send_event(run_attempt_proto)
+
+
+TRAIN_STATE_ACTOR_NAME = "train_v2_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 the Ray Train state actor singleton.
+
+    This is a long-living, detached actor living on the head node
+    that gets initialized when the first Train run happens on the
+    Ray cluster.
+    """
+    with _state_actor_lock:
+        state_actor = (
+            ray.remote(TrainStateActor)
+            .options(
+                num_cpus=0,
+                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",
+                max_restarts=-1,
+                max_task_retries=-1,
+            )
+            .remote(
+                enable_state_actor_reconciliation=ray_constants.env_bool(
+                    ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
+                    DEFAULT_ENABLE_STATE_ACTOR_RECONCILIATION,
+                ),
+                reconciliation_interval_s=float(
+                    os.getenv(
+                        STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR,
+                        DEFAULT_STATE_ACTOR_RECONCILIATION_INTERVAL_S,
+                    )
+                ),
+            )
+        )
+
+    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/v2/_internal/state/state_manager.py

@@ -0,0 +1,326 @@
+import logging
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import ray
+from ray.actor import ActorHandle
+from ray.train import BackendConfig
+from ray.train._internal.data_config import DataConfig
+from ray.train.v2._internal.execution.context import DistributedContext
+from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
+    ResizeDecision,
+)
+from ray.train.v2._internal.execution.worker_group import ActorMetadata, Worker
+from ray.train.v2._internal.state.schema import (
+    ActorStatus,
+    BackendConfig as BackendConfigSchema,
+    CheckpointConfig as CheckpointConfigSchema,
+    FailureConfig as FailureConfigSchema,
+    RunAttemptStatus,
+    RunConfig as RunConfigSchema,
+    RunSettings,
+    RunStatus,
+    ScalingConfig as ScalingConfigSchema,
+    TrainResources,
+    TrainRun,
+    TrainRunAttempt,
+    TrainWorker,
+)
+from ray.train.v2._internal.state.state_actor import get_or_create_state_actor
+from ray.train.v2._internal.state.util import (
+    construct_data_config,
+    current_time_ns,
+    mark_workers_dead,
+    update_train_run_aborted,
+    update_train_run_attempt_aborted,
+)
+from ray.train.v2._internal.util import TrainingFramework
+from ray.train.v2.api.config import RunConfig, ScalingConfig
+
+logger = logging.getLogger(__name__)
+
+
+class TrainStateManager:
+    """Manages the state of a train run and run attempts."""
+
+    def __init__(self) -> None:
+        self._state_actor = get_or_create_state_actor()
+        # NOTE: All runs and attempts are stored in memory.
+        # This may be a memory issue for large runs.
+        self._runs: Dict[str, TrainRun] = {}
+        # {run_id: {attempt_id: TrainRunAttempt}}
+        self._run_attempts: Dict[str, Dict[str, TrainRunAttempt]] = defaultdict(dict)
+
+    def create_train_run(
+        self,
+        id: str,
+        name: str,
+        job_id: str,
+        controller_actor_id: str,
+        controller_log_file_path: str,
+        run_config: RunConfig,
+        train_loop_config: Optional[Dict],
+        scaling_config: ScalingConfig,
+        backend_config: BackendConfig,
+        datasets: Dict[str, ray.data.Dataset],
+        dataset_config: DataConfig,
+    ) -> None:
+        run_config_schema = RunConfigSchema(
+            name=run_config.name,
+            failure_config=FailureConfigSchema(
+                max_failures=run_config.failure_config.max_failures,
+                controller_failure_limit=run_config.failure_config.controller_failure_limit,
+            ),
+            worker_runtime_env=run_config.worker_runtime_env,
+            checkpoint_config=CheckpointConfigSchema(
+                num_to_keep=run_config.checkpoint_config.num_to_keep,
+                checkpoint_score_attribute=run_config.checkpoint_config.checkpoint_score_attribute,
+                checkpoint_score_order=run_config.checkpoint_config.checkpoint_score_order,
+            ),
+            storage_path=run_config.storage_path,
+            storage_filesystem=run_config.storage_filesystem,
+        )
+
+        scaling_config_schema = ScalingConfigSchema(
+            num_workers=scaling_config.num_workers,
+            use_gpu=scaling_config.use_gpu,
+            resources_per_worker=scaling_config.resources_per_worker,
+            placement_strategy=scaling_config.placement_strategy,
+            accelerator_type=scaling_config.accelerator_type,
+            use_tpu=scaling_config.use_tpu,
+            topology=scaling_config.topology,
+            bundle_label_selector=scaling_config.label_selector,
+        )
+
+        backend_config_schema = BackendConfigSchema(
+            framework=backend_config.framework,
+            config=backend_config.to_dict(),
+        )
+
+        run_settings = RunSettings(
+            train_loop_config=train_loop_config,
+            backend_config=backend_config_schema,
+            scaling_config=scaling_config_schema,
+            datasets=list(datasets.keys()),
+            data_config=construct_data_config(dataset_config),
+            run_config=run_config_schema,
+        )
+
+        run = TrainRun(
+            id=id,
+            name=name,
+            job_id=job_id,
+            status=RunStatus.INITIALIZING,
+            status_detail=None,
+            controller_actor_id=controller_actor_id,
+            start_time_ns=current_time_ns(),
+            end_time_ns=None,
+            controller_log_file_path=controller_log_file_path,
+            framework_versions={"ray": ray.__version__},
+            run_settings=run_settings,
+        )
+        self._runs[run.id] = run
+        # Block so the initial run state isn't lost if the controller exits
+        # right after. Without this, the .remote() task could still be in the
+        # caller's outbound queue when the controller dies, leaving the state
+        # actor with no record of the run.
+        self._create_or_update_train_run(run, block=True)
+
+    def update_train_run_scheduling(
+        self,
+        run_id: str,
+        resize_decision: Optional[ResizeDecision] = None,
+    ) -> None:
+        if resize_decision is not None:
+            status_detail = _get_scheduling_status_detail(
+                resize_decision.num_workers, resize_decision.resources_per_worker
+            )
+        else:
+            status_detail = None
+
+        run = self._runs[run_id]
+        run.status = RunStatus.SCHEDULING
+        run.status_detail = status_detail
+        self._create_or_update_train_run(run)
+
+    def update_train_run_running(
+        self,
+        run_id: str,
+    ) -> None:
+        run = self._runs[run_id]
+        run.status = RunStatus.RUNNING
+        run.status_detail = None
+        self._create_or_update_train_run(run)
+
+    def update_train_run_restarting(
+        self,
+        run_id: str,
+    ) -> None:
+        run = self._runs[run_id]
+        run.status = RunStatus.RESTARTING
+        run.status_detail = None
+        self._create_or_update_train_run(run)
+
+    def update_train_run_resizing(
+        self,
+        run_id: str,
+    ) -> None:
+        run = self._runs[run_id]
+        run.status = RunStatus.RESIZING
+        run.status_detail = None
+        self._create_or_update_train_run(run)
+
+    def update_train_run_finished(
+        self,
+        run_id: str,
+    ):
+        run = self._runs[run_id]
+        run.status = RunStatus.FINISHED
+        run.status_detail = None
+        run.end_time_ns = current_time_ns()
+        # Block on terminal status so the final state isn't lost if the controller exits right after.
+        self._create_or_update_train_run(run, block=True)
+
+    def update_train_run_errored(
+        self,
+        run_id: str,
+        status_detail: str,
+    ):
+        run = self._runs[run_id]
+        run.status = RunStatus.ERRORED
+        run.status_detail = status_detail
+        run.end_time_ns = current_time_ns()
+        # Block on terminal status so the final state isn't lost if the controller exits right after.
+        self._create_or_update_train_run(run, block=True)
+
+    def update_train_run_aborted(
+        self,
+        run_id: str,
+    ):
+        run = self._runs[run_id]
+        update_train_run_aborted(run=run, graceful=True)
+        # Block on terminal status so the final state isn't lost if the controller exits right after.
+        self._create_or_update_train_run(run, block=True)
+
+    def update_train_run_framework_versions(
+        self, run_id: str, framework_versions: Dict[str, str]
+    ):
+        run = self._runs[run_id]
+        run.framework_versions = framework_versions
+        self._create_or_update_train_run(run)
+
+    def create_train_run_attempt(
+        self,
+        run_id: str,
+        attempt_id: str,
+        num_workers: int,
+        resources_per_worker: Dict[str, float],
+    ) -> None:
+        status_detail = _get_scheduling_status_detail(num_workers, resources_per_worker)
+        resources = [
+            TrainResources(resources=resources_per_worker) for _ in range(num_workers)
+        ]
+        run_attempt = TrainRunAttempt(
+            run_id=run_id,
+            attempt_id=attempt_id,
+            start_time_ns=current_time_ns(),
+            status=RunAttemptStatus.PENDING,
+            status_detail=status_detail,
+            resources=resources,
+            workers=[],  # Not started yet.
+        )
+
+        self._run_attempts[run_id][attempt_id] = run_attempt
+        self._create_or_update_train_run_attempt(run_attempt)
+
+    def update_train_run_attempt_running(
+        self, run_id: str, attempt_id: str, workers: List[Worker]
+    ) -> None:
+        def _convert_worker(worker: Worker) -> TrainWorker:
+
+            actor: ActorHandle = worker.actor
+            distributed_context: DistributedContext = worker.distributed_context
+            actor_metadata: ActorMetadata = worker.metadata
+
+            return TrainWorker(
+                world_rank=distributed_context.world_rank,
+                local_rank=distributed_context.local_rank,
+                node_rank=distributed_context.node_rank,
+                actor_id=actor._actor_id.hex(),
+                node_id=actor_metadata.node_id,
+                node_ip=actor_metadata.node_ip,
+                pid=actor_metadata.pid,
+                gpu_ids=actor_metadata.gpu_ids,
+                status=ActorStatus.ALIVE,
+                resources=TrainResources(resources=worker.resources),
+                log_file_path=worker.log_file_path,
+            )
+
+        workers: List[TrainWorker] = [_convert_worker(worker) for worker in workers]
+
+        run_attempt = self._run_attempts[run_id][attempt_id]
+        run_attempt.status = RunAttemptStatus.RUNNING
+        run_attempt.status_detail = None
+        run_attempt.workers = workers
+        self._create_or_update_train_run_attempt(run_attempt)
+
+    def update_train_run_attempt_finished(
+        self,
+        run_id: str,
+        attempt_id: str,
+    ):
+        run_attempt = self._run_attempts[run_id][attempt_id]
+        run_attempt.status = RunAttemptStatus.FINISHED
+        run_attempt.status_detail = None
+        run_attempt.end_time_ns = current_time_ns()
+        mark_workers_dead(run_attempt)
+        # Block to avoid case where controller is dead but attempt is not terminal.
+        self._create_or_update_train_run_attempt(run_attempt, block=True)
+
+    def update_train_run_attempt_errored(
+        self,
+        run_id: str,
+        attempt_id: str,
+        status_detail: str,
+    ):
+        run_attempt = self._run_attempts[run_id][attempt_id]
+        run_attempt.status = RunAttemptStatus.ERRORED
+        run_attempt.status_detail = status_detail
+        run_attempt.end_time_ns = current_time_ns()
+        mark_workers_dead(run_attempt)
+        # Block to avoid case where controller is dead but attempt is not terminal.
+        self._create_or_update_train_run_attempt(run_attempt, block=True)
+
+    def update_train_run_attempt_aborted(
+        self,
+        run_id: str,
+        attempt_id: str,
+    ):
+        run_attempt = self._run_attempts[run_id][attempt_id]
+        update_train_run_attempt_aborted(run_attempt=run_attempt, graceful=True)
+        # Block to avoid case where controller is dead but attempt is not terminal.
+        self._create_or_update_train_run_attempt(run_attempt, block=True)
+
+    def get_train_run_framework(self, run_id: str) -> Optional[TrainingFramework]:
+        run = self._runs[run_id]
+        return run.run_settings.backend_config.framework
+
+    def _create_or_update_train_run(
+        self, run: TrainRun, *, block: bool = False
+    ) -> None:
+        ref = self._state_actor.create_or_update_train_run.remote(run)
+        if block:
+            ray.get(ref)
+
+    def _create_or_update_train_run_attempt(
+        self, run_attempt: TrainRunAttempt, *, block: bool = False
+    ) -> None:
+        ref = self._state_actor.create_or_update_train_run_attempt.remote(run_attempt)
+        if block:
+            ray.get(ref)
+
+
+def _get_scheduling_status_detail(
+    num_workers: int, resources_per_worker: Dict[str, float]
+) -> str:
+    return f"Scheduling {num_workers} workers, each requiring: {resources_per_worker}."

python/ray/train/v2/_internal/state/util.py

@@ -0,0 +1,97 @@
+import time
+
+from ray.data._internal.execution.interfaces.execution_options import ExecutionOptions
+from ray.train._internal.data_config import DataConfig
+from ray.train.v2._internal.state.schema import (
+    ActorStatus,
+    DataConfig as DataConfigSchema,
+    DataExecutionOptions,
+    ExecutionOptions as ExecutionOptionsSchema,
+    RunAttemptStatus,
+    RunStatus,
+    TrainRun,
+    TrainRunAttempt,
+)
+from ray.util.state import get_actor
+
+_GRACEFUL_ABORT_STATUS_DETAIL = "Run aborted due to user interrupt (SIGINT)."
+_DEAD_CONTROLLER_ABORT_STATUS_DETAIL = (
+    "Run aborted because the driver process exited unexpectedly."
+)
+
+
+def update_train_run_aborted(run: TrainRun, graceful: bool) -> None:
+    run.status = RunStatus.ABORTED
+    if graceful:
+        run.status_detail = _GRACEFUL_ABORT_STATUS_DETAIL
+    else:
+        run.status_detail = _DEAD_CONTROLLER_ABORT_STATUS_DETAIL
+    run.end_time_ns = current_time_ns()
+
+
+def update_train_run_attempt_aborted(
+    run_attempt: TrainRunAttempt, graceful: bool
+) -> None:
+    if graceful:
+        run_attempt.status_detail = _GRACEFUL_ABORT_STATUS_DETAIL
+    else:
+        run_attempt.status_detail = _DEAD_CONTROLLER_ABORT_STATUS_DETAIL
+    run_attempt.status = RunAttemptStatus.ABORTED
+    run_attempt.end_time_ns = current_time_ns()
+    mark_workers_dead(run_attempt)
+
+
+def mark_workers_dead(run_attempt: TrainRunAttempt) -> None:
+    for worker in run_attempt.workers:
+        worker.status = ActorStatus.DEAD
+
+
+def current_time_ns() -> int:
+    return time.time_ns()
+
+
+def is_actor_alive(actor_id: str, timeout: int) -> bool:
+    """Returns whether actor is alive."""
+    actor_state = get_actor(actor_id, timeout=timeout)
+    return actor_state and actor_state.state != "DEAD"
+
+
+def construct_data_config(data_config: DataConfig) -> DataConfigSchema:
+    """Serialize a user-facing DataConfig into the exportable schema.
+
+    Note: This function assumes data_config._execution_options (a defaultdict)
+    hasn't been read between initialization of the field and this function call.
+    Any read materializes a dataset key and affects the data config shape,
+    wrongly capturing a per dataset execution options even if the user only
+    provided a default.
+    """
+    exec_options = data_config._execution_options
+
+    per_dataset_execution_options = {}
+    if exec_options:
+        per_dataset_execution_options = {
+            ds_name: execution_options_to_model(opts)
+            for ds_name, opts in exec_options.items()
+        }
+
+    return DataConfigSchema(
+        datasets_to_split=data_config._datasets_to_split,
+        data_execution_options=DataExecutionOptions(
+            default=execution_options_to_model(exec_options.default_factory()),
+            per_dataset_execution_options=per_dataset_execution_options,
+        ),
+        enable_shard_locality=data_config._enable_shard_locality,
+    )
+
+
+def execution_options_to_model(
+    execution_options: ExecutionOptions,
+) -> ExecutionOptionsSchema:
+    """Convert a ray.data ExecutionOptions object into the export schema model."""
+    return ExecutionOptionsSchema(
+        resource_limits=execution_options.resource_limits.to_resource_dict(),
+        exclude_resources=execution_options.exclude_resources.to_resource_dict(),
+        preserve_order=execution_options.preserve_order,
+        actor_locality_enabled=execution_options.actor_locality_enabled,
+        verbose_progress=execution_options.verbose_progress,
+    )

python/ray/train/v2/_internal/util.py

@@ -0,0 +1,369 @@
+import asyncio
+import contextlib
+import functools
+import logging
+import threading
+import time
+import traceback
+from datetime import datetime
+from enum import Enum
+from typing import (
+    Any,
+    Callable,
+    ContextManager,
+    Dict,
+    Generator,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+)
+
+import ray
+from ray.train._internal.utils import count_required_parameters
+from ray.train.v2._internal.exceptions import UserExceptionWithTraceback
+
+logger = logging.getLogger(__name__)
+
+
+T = TypeVar("T")
+
+
+def bundle_to_remote_args(bundle: dict) -> dict:
+    """Convert a bundle of resources to Ray actor/task arguments.
+
+    >>> bundle_to_remote_args({"GPU": 1, "memory": 1, "custom": 0.1})
+    {'num_cpus': 0, 'num_gpus': 1, 'memory': 1, 'resources': {'custom': 0.1}}
+    """
+    bundle = bundle.copy()
+    args = {
+        "num_cpus": bundle.pop("CPU", 0),
+        "num_gpus": bundle.pop("GPU", 0),
+        "memory": bundle.pop("memory", 0),
+    }
+    if bundle:
+        args["resources"] = bundle
+    return args
+
+
+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_loop_per_worker",
+) -> 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.
+
+    Returns:
+        A valid training function.
+
+    Raises:
+        ValueError: if the input ``train_func`` is invalid.
+    """
+    num_required_params = count_required_parameters(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)
+
+    if num_required_params == 1:
+        config = config or {}
+
+        @functools.wraps(train_func)
+        def train_fn():
+            with train_func_context():
+                return train_func(config)
+
+    else:  # num_params == 0
+
+        @functools.wraps(train_func)
+        def train_fn():
+            with train_func_context():
+                return train_func()
+
+    return train_fn
+
+
+class TrainingFramework(Enum):
+    TORCH = "torch"
+    JAX = "jax"
+    TENSORFLOW = "tensorflow"
+    XGBOOST = "xgboost"
+    LIGHTGBM = "lightgbm"
+
+    def module_names(self) -> tuple[str, ...]:
+        """Returns the relevant module names for the training framework.
+
+        These module names are used by Train state version collection (see
+        `_get_framework_version`) to gather versions of key framework-related packages.
+
+        Note: If adding a new module, make sure to use the module name rather than
+        the distribution name. (e.g. sklearn instead of scikit-learn)
+        """
+        if self is TrainingFramework.TORCH:
+            return ("torch",)
+        if self is TrainingFramework.JAX:
+            return ("jax", "jaxlib")
+        if self is TrainingFramework.TENSORFLOW:
+            return ("tensorflow", "keras")
+        if self is TrainingFramework.XGBOOST:
+            return ("xgboost",)
+        if self is TrainingFramework.LIGHTGBM:
+            return ("lightgbm",)
+
+        return (self.value,)
+
+
+class ObjectRefWrapper(Generic[T]):
+    """Thin wrapper around ray.put to manually control dereferencing."""
+
+    def __init__(self, obj: T):
+        self._ref = ray.put(obj)
+
+    def get(self) -> T:
+        return ray.get(self._ref)
+
+
+def date_str(include_ms: bool = False):
+    pattern = "%Y-%m-%d_%H-%M-%S"
+    if include_ms:
+        pattern += ".%f"
+    return datetime.today().strftime(pattern)
+
+
+def time_monotonic():
+    return time.monotonic()
+
+
+def _copy_doc(copy_func):
+    def wrapped(func):
+        func.__doc__ = copy_func.__doc__
+        return func
+
+    return wrapped
+
+
+@contextlib.contextmanager
+def invoke_context_managers(
+    context_managers: List[ContextManager],
+) -> Generator[None, None, None]:
+    """
+    Utility to invoke a list of context managers and yield sequentially.
+
+    Args:
+        context_managers: List of context managers to invoke.
+    """
+    with contextlib.ExitStack() as stack:
+        for context_manager in context_managers:
+            stack.enter_context(context_manager())
+        yield
+
+
+def get_module_name(obj: object) -> str:
+    """Returns the full module name of the given object, including its qualified name.
+
+    Args:
+        obj: The object (class, function, etc.) whose module name is required.
+
+    Returns:
+        Full module and qualified name as a string.
+    """
+    return f"{obj.__module__}.{obj.__qualname__}"
+
+
+def get_callable_name(fn: Callable) -> str:
+    """Returns a readable name for any callable.
+
+    Args:
+        fn: The callable to extract a name from.
+
+    Returns:
+        A human-readable name for the callable.
+
+    Examples:
+
+        >>> get_callable_name(lambda x: x)
+        '<lambda>'
+        >>> def foo(a, b): pass
+        >>> get_callable_name(foo)
+        'foo'
+        >>> from functools import partial
+        >>> bar = partial(partial(foo, a=1), b=2)
+        >>> get_callable_name(bar)
+        'foo'
+        >>> class Dummy:
+        ...     def __call__(self, a, b): pass
+        >>> get_callable_name(Dummy())
+        'Dummy'
+    """
+    if isinstance(fn, functools.partial):
+        return get_callable_name(fn.func)
+
+    # Use __name__ for regular functions and lambdas
+    if hasattr(fn, "__name__"):
+        return fn.__name__
+
+    # Fallback to the class name for objects that implement __call__
+    return fn.__class__.__name__
+
+
+def construct_user_exception_with_traceback(
+    e: BaseException, exclude_frames: int = 0
+) -> UserExceptionWithTraceback:
+    """Construct a UserExceptionWithTraceback from a base exception.
+
+    Args:
+        e: The base exception to construct a UserExceptionWithTraceback from.
+        exclude_frames: The number of frames to exclude from the beginnning of
+            the traceback.
+
+    Returns:
+        A UserExceptionWithTraceback object.
+    """
+    # TODO(justinvyu): This is brittle and may break if the call stack
+    # changes. Figure out a more robust way to exclude these frames.
+    exc_traceback_str = traceback.format_exc(
+        limit=-(len(traceback.extract_tb(e.__traceback__)) - exclude_frames)
+    )
+    logger.error(f"Error in training function:\n{exc_traceback_str}")
+    return UserExceptionWithTraceback(e, traceback_str=exc_traceback_str)
+
+
+def _in_ray_train_worker() -> bool:
+    """Check if the current process is a Ray Train V2 worker."""
+    from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
+
+    try:
+        get_train_fn_utils()
+        return True
+    except RuntimeError:
+        return False
+
+
+def requires_train_worker(raise_in_tune_session: bool = False) -> Callable:
+    """Check that the caller is a Ray Train worker spawned by Ray Train,
+    with access to training function utilities.
+
+    Args:
+        raise_in_tune_session: Whether to raise a specific error message if the caller
+            is in a Tune session. If True, will raise a DeprecationWarning.
+
+    Returns:
+        A decorator that performs this check, which raises an error if the caller
+        is not a Ray Train worker.
+    """
+
+    def _wrap(fn: Callable) -> Callable:
+        @functools.wraps(fn)
+        def _wrapped_fn(*args, **kwargs):
+            from ray.tune.trainable.trainable_fn_utils import _in_tune_session
+
+            if raise_in_tune_session and _in_tune_session():
+                raise DeprecationWarning(
+                    f"`ray.train.{fn.__name__}` is deprecated when running in a function "
+                    "passed to Ray Tune. Please use the equivalent `ray.tune` API instead. "
+                    "See this issue for more context: "
+                    "https://github.com/ray-project/ray/issues/49454"
+                )
+
+            if not _in_ray_train_worker():
+                raise RuntimeError(
+                    f"`{fn.__name__}` cannot be used outside of a Ray Train training function. "
+                    "You are calling this API from the driver or another non-training process. "
+                    "These utilities are only available within a function launched by `trainer.fit()`."
+                )
+            return fn(*args, **kwargs)
+
+        return _wrapped_fn
+
+    return _wrap
+
+
+async def wait_with_logging(
+    condition: asyncio.Condition,
+    predicate: Optional[Callable[[], bool]] = None,
+    generate_warning_message: Optional[Callable[[], str]] = None,
+    warn_interval_s: float = 60,
+    timeout_s: Optional[float] = None,
+):
+    """Waits for condition to be notified, logging warnings and eventually timing out.
+
+    You must acquire the condition before calling this function.
+
+    Args:
+        condition: The condition to wait for.
+        predicate: Wait until this predicate is True. If None, wait until the condition
+            is notified.
+        generate_warning_message: A function that generates the warning message to log.
+            If None, no warning is logged.
+        warn_interval_s: The interval in seconds to log a warning.
+        timeout_s: The timeout in seconds. Defaults to``None`` to not time out.
+    """
+
+    async def _wait_loop():
+        while True:
+            try:
+                await asyncio.wait_for(
+                    condition.wait()
+                    if predicate is None
+                    else condition.wait_for(predicate),
+                    timeout=warn_interval_s,
+                )
+                return
+            # asyncio.wait_for() raises `asyncio.TimeoutError` for asyncio<=3.10
+            # and raises `TimeoutError` for asyncio>=3.11
+            # https://docs.python.org/3/library/asyncio-task.html#asyncio.wait_for
+            except (asyncio.TimeoutError, TimeoutError):
+                if generate_warning_message is not None:
+                    warning_message = generate_warning_message()
+                    logger.warning(warning_message)
+
+    await asyncio.wait_for(
+        _wait_loop(),
+        timeout=timeout_s,
+    )
+
+
+@contextlib.contextmanager
+def context_watchdog(fn: Callable, *args: Any) -> Iterator[None]:
+    """Run a function in a background thread for the duration of the context.
+
+    The function is started in a daemon thread on entry. On exit, a
+    threading.Event is set to signal the thread to stop. The function is
+    responsible for checking the event and returning promptly once it is set.
+
+    Args:
+        fn: A function whose first argument is a threading.Event stop signal.
+            The function should return when stop_event.is_set() or
+            stop_event.wait(...) returns True.
+        *args: Additional arguments forwarded to fn after the stop event.
+
+    Yields:
+        None: Control is yielded to the caller while the watchdog thread runs.
+    """
+    stop_event = threading.Event()
+    thread = threading.Thread(
+        target=fn,
+        args=(stop_event, *args),
+        daemon=True,  # thread will end even if the finally is bypassed by an abnormal exit
+    )
+    thread.start()
+    try:
+        yield
+    finally:
+        stop_event.set()
+        thread.join()