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