ray-project--ray/python/ray/train/v2/api/data_parallel_trainer.py

import logging
import signal
import sys
import threading
from typing import Any, Callable, Dict, List, Optional, Union

import ray
from ray._common.constants import RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR
from ray._common.usage import usage_lib
from ray._private.ray_constants import env_bool
from ray.actor import ActorHandle
from ray.air._internal.usage import tag_train_v2_trainer
from ray.train import (
    BackendConfig,
    Checkpoint,
    DataConfig,
    Result,
    RunConfig,
    ScalingConfig,
)
from ray.train.base_trainer import (
    _RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING,
    _TRAINER_RESTORE_DEPRECATION_WARNING,
)
from ray.train.constants import RAY_CHDIR_TO_TRIAL_DIR, RAY_TRAIN_ENABLE_STATE_TRACKING
from ray.train.context import _GET_METADATA_DEPRECATION_MESSAGE
from ray.train.v2._internal.callbacks import (
    AcceleratorSetupCallback,
    BackendSetupCallback,
    DatasetsCallback,
    WorkingDirectorySetupCallback,
)
from ray.train.v2._internal.callbacks.env_callback import _initialize_env_callbacks
from ray.train.v2._internal.callbacks.metrics import (
    ControllerMetricsCallback,
    WorkerMetricsCallback,
)
from ray.train.v2._internal.callbacks.placement_group_callback import (
    PlacementGroupCleanerCallback,
)
from ray.train.v2._internal.callbacks.state_manager import StateManagerCallback
from ray.train.v2._internal.callbacks.user_callback import UserCallbackHandler
from ray.train.v2._internal.constants import (
    DEFAULT_RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_VALUE,
    METRICS_ENABLED_ENV_VAR,
    V2_ENABLED_ENV_VAR,
    get_env_vars_to_propagate,
    is_v2_enabled,
)
from ray.train.v2._internal.data_integration.interfaces import GenDataset
from ray.train.v2._internal.execution.callback import RayTrainCallback
from ray.train.v2._internal.execution.context import TrainRunContext
from ray.train.v2._internal.execution.controller import TrainController
from ray.train.v2._internal.execution.failure_handling import create_failure_policy
from ray.train.v2._internal.execution.local_mode.utils import LocalController
from ray.train.v2._internal.execution.scaling_policy import create_scaling_policy
from ray.train.v2._internal.util import ObjectRefWrapper, construct_train_func
from ray.train.v2.api.callback import UserCallback
from ray.train.v2.api.validation_config import ValidationConfig
from ray.util.annotations import Deprecated, DeveloperAPI

logger = logging.getLogger(__name__)


@DeveloperAPI
class DataParallelTrainer:
    """Base class for distributed data parallel training on Ray.

    This class supports the SPMD parallelization pattern, where a single
    training function is executed in parallel across multiple workers,
    and different shards of data are processed by each worker.
    """

    def __init__(
        self,
        train_loop_per_worker: Union[Callable[[], Any], Callable[[Dict], Any]],
        *,
        train_loop_config: Optional[Dict] = None,
        backend_config: Optional[BackendConfig] = None,
        scaling_config: Optional[ScalingConfig] = None,
        run_config: Optional[RunConfig] = None,
        datasets: Optional[Dict[str, GenDataset]] = None,
        dataset_config: Optional[DataConfig] = None,
        # TODO: [Deprecated] Remove in future release
        resume_from_checkpoint: Optional[Checkpoint] = None,
        metadata: Optional[Dict[str, Any]] = None,
        validation_config: Optional[ValidationConfig] = None,
    ):
        self.run_config = run_config or RunConfig()
        self.train_loop_per_worker = train_loop_per_worker
        self.validation_config = validation_config
        self.train_loop_config = train_loop_config
        self.scaling_config = scaling_config or ScalingConfig()
        self.backend_config = backend_config or BackendConfig()
        self.datasets = datasets or {}
        self.data_config = dataset_config or DataConfig()

        self.running_in_local_mode = self.scaling_config.num_workers == 0

        self.train_run_context = TrainRunContext(
            run_config=self.run_config,
            train_loop_config=self.train_loop_config,
            scaling_config=self.scaling_config,
            backend_config=self.backend_config,
            dataset_config=self.data_config,
        )

        if resume_from_checkpoint is not None:
            raise DeprecationWarning(_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING)

        if metadata is not None:
            raise DeprecationWarning(_GET_METADATA_DEPRECATION_MESSAGE)

        self._validate_configs()

        usage_lib.record_library_usage("train")
        tag_train_v2_trainer(self)
        if self.scaling_config.elasticity_enabled:
            usage_lib.record_extra_usage_tag(
                usage_lib.TagKey.TRAIN_ELASTICITY_ENABLED, "1"
            )

    def _validate_configs(self):
        if not is_v2_enabled():
            raise ValueError(
                f"Ray Train V2 must be enabled with `{V2_ENABLED_ENV_VAR}=1` "
                "when using this V2 Trainer API."
            )

        from ray.train.v2.api.config import (
            RunConfig as RunConfigV2,
            ScalingConfig as ScalingConfigV2,
        )

        if not isinstance(self.run_config, RunConfigV2):
            raise ValueError(
                f"Invalid `RunConfig` type: {self.run_config.__class__}. "
                "Use `ray.train.RunConfig` instead. "
                "See this issue for more context: "
                "https://github.com/ray-project/ray/issues/49454"
            )

        if not isinstance(self.scaling_config, ScalingConfigV2):
            raise ValueError(
                f"Invalid `ScalingConfig` type: {self.scaling_config.__class__}. "
                "Use `ray.train.ScalingConfig` instead. "
                "See this issue for more context: "
                "https://github.com/ray-project/ray/issues/49454"
            )

    def _get_train_func(self) -> Callable[[], Any]:
        return construct_train_func(
            self.train_loop_per_worker,
            config=self.train_loop_config,
            train_func_context=self.backend_config.train_func_context,
            fn_arg_name="train_loop_per_worker",
        )

    def fit(self) -> Result:
        """Launches the Ray Train controller to run training on workers.

        Returns:
            A Result object containing the training result.

        Raises:
            ray.train.TrainingFailedError: This is a union of the ControllerError and WorkerGroupError.
                This returns a :class:`ray.train.ControllerError` if internal Ray Train controller logic
                encounters a non-retryable error or reaches the controller failure limit configured in `FailureConfig`.
                This returns a :class:`ray.train.WorkerGroupError` if one or more workers fail during
                training and reaches the worker group failure limit configured in `FailureConfig(max_failures)`.
        """
        train_fn = self._get_train_func()
        if self.running_in_local_mode:
            return self._initialize_and_run_local_controller(train_fn)
        else:
            train_fn_ref = ObjectRefWrapper(train_fn)

            result = self._initialize_and_run_controller(
                train_fn_ref=train_fn_ref,
                scaling_policy=create_scaling_policy(self.scaling_config),
                failure_policy=create_failure_policy(self.run_config.failure_config),
                train_run_context=self.train_run_context,
                callbacks=self._create_default_callbacks(),
                validation_config=self.validation_config,
            )

            if result.error:
                # NOTE: If the training run errored out, raise an error back to the
                # user's driver script.
                # For example, if the Train `FailurePolicy` runs out of retries,
                # and one of the workers errors. The controller will exit, and
                # the error will be raised here.
                raise result.error

            return result

    def _get_local_controller(self) -> LocalController:
        return LocalController(
            experiment_name=self.run_config.name,
            datasets=self.datasets,
        )

    def _create_default_callbacks(self) -> List[RayTrainCallback]:
        # Initialize callbacks from environment variable
        callbacks = _initialize_env_callbacks()

        accelerator_setup_callback = AcceleratorSetupCallback(
            self.backend_config, self.scaling_config
        )
        backend_setup_callback = BackendSetupCallback(self.backend_config)
        datasets_callback = DatasetsCallback(
            train_run_context=self.train_run_context,
            datasets=self.datasets,
        )
        placement_group_cleaner_callback = PlacementGroupCleanerCallback()
        callbacks.extend(
            [
                accelerator_setup_callback,
                backend_setup_callback,
                placement_group_cleaner_callback,
                datasets_callback,
            ]
        )
        if env_bool(RAY_CHDIR_TO_TRIAL_DIR, True):
            working_directory_setup_callback = WorkingDirectorySetupCallback()
            callbacks.append(working_directory_setup_callback)

        if env_bool(METRICS_ENABLED_ENV_VAR, True):
            callbacks.append(ControllerMetricsCallback())
            callbacks.append(WorkerMetricsCallback(self.train_run_context))

        if env_bool(RAY_TRAIN_ENABLE_STATE_TRACKING, False):
            callbacks.append(StateManagerCallback(datasets=self.datasets))

        run_config_callbacks = (
            self.run_config.callbacks if self.run_config.callbacks is not None else []
        )

        # Add internal callback that invokes all user-defined callbacks.
        user_callbacks = [
            cb for cb in run_config_callbacks if isinstance(cb, UserCallback)
        ]
        callbacks.append(
            UserCallbackHandler(
                user_callbacks=user_callbacks, train_run_context=self.train_run_context
            )
        )

        # Append all other callbacks to the full list. This allows custom workarounds
        # built on top of internal callbacks to work.
        callbacks.extend(
            [cb for cb in run_config_callbacks if not isinstance(cb, UserCallback)]
        )
        return callbacks

    def _initialize_and_run_local_controller(
        self, train_func: Callable[[], Any]
    ) -> Result:
        return self._get_local_controller().run(train_func)

    def _initialize_and_run_controller(self, **controller_init_kwargs) -> Result:
        env_vars = get_env_vars_to_propagate()
        env_vars.setdefault(
            RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR,
            DEFAULT_RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_VALUE,
        )

        # Attach the controller to the node running the driver script.
        controller_actor_cls = ray.remote(
            num_cpus=0,
            label_selector={
                ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
            },
            # TODO: Extract env variables that affect controller behavior
            # and pass them as explicit args
            runtime_env={"env_vars": env_vars},
        )(TrainController)

        controller = controller_actor_cls.remote(**controller_init_kwargs)

        # If this is not the main thread - as is the case when running in Tune -
        # registering the SIGINT handler raises an exception.
        if threading.current_thread() is threading.main_thread():
            self._register_sigint_handler(controller)

        ray.get(controller.run.remote())
        return ray.get(controller.get_result.remote())

    def _register_sigint_handler(self, controller: ActorHandle[TrainController]):
        """Register SIGINT handler so user Ctrl C gracefully aborts run."""
        sigint_count = 0

        def sigint_handler(signum, frame):
            logger.info(
                "Received SIGINT. Gracefully aborting the training run — this "
                "may take a few seconds. To forcefully abort immediately, you "
                "can send a different signal, such as SIGKILL."
            )
            nonlocal sigint_count
            sigint_count += 1
            if sigint_count >= 3:
                logger.info(
                    "Received SIGINT at least 3 times. "
                    "Forcefully aborting the training run."
                )
                sys.exit(0)
            if sigint_count <= 1:
                try:
                    ray.get(controller.abort.remote())
                except ray.exceptions.RayActorError:
                    # We catch the error and exit 0 to indicate graceful termination.
                    # However, for some reason the process still exits with 1.
                    sys.exit(0)

        signal.signal(signal.SIGINT, sigint_handler)

    @classmethod
    @Deprecated
    def restore(cls, *args, **kwargs):
        """[Deprecated] Restores a Train experiment from a previously
        interrupted/failed run.

        This method is deprecated and will be removed in a future release.
        """
        raise DeprecationWarning(_TRAINER_RESTORE_DEPRECATION_WARNING)

    @classmethod
    @Deprecated
    def can_restore(cls, *args, **kwargs):
        """[Deprecated] Checks if a Train experiment can be restored from
        a previously interrupted/failed run.

        This method is deprecated and will be removed in a future release.
        """
        raise DeprecationWarning(_TRAINER_RESTORE_DEPRECATION_WARNING)