wehub-resource-sync
859 lines
31 KiB
Python
859 lines
31 KiB
Python
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
|