chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,335 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user