chore: import upstream snapshot with attribution

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