wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/train/_internal/__init__.py
View File
python/ray/train/_internal/accelerator.py
+5
View File
@@ -0,0 +1,5 @@
import abc
class Accelerator(abc.ABC):
"""A utility that contains methods to accelerate training."""
python/ray/train/_internal/backend_executor.py
+858
View File
@@ -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
python/ray/train/_internal/base_worker_group.py
+105
View File
@@ -0,0 +1,105 @@
"""Abstract base class for WorkerGroup implementations.
This module defines the common base class that both V1 and V2 WorkerGroup
implementations should inherit from to ensure backend compatibility.
"""
import abc
from typing import Callable, List, TypeVar
from ray.types import ObjectRef
from ray.util.annotations import DeveloperAPI
T = TypeVar("T")
@DeveloperAPI
class BaseWorkerGroup(abc.ABC):
"""Abstract base class for WorkerGroup implementations.
This base class defines the minimal set of methods that backend classes
expect from WorkerGroup implementations. Both V1 and V2 WorkerGroup
classes should inherit from this base class to ensure compatibility with
all backend configurations.
The interface focuses on the core operations that backends need:
- Executing functions on workers
- Getting worker count and resource allocation
"""
@abc.abstractmethod
def execute(self, func: Callable[..., T], *args, **kwargs) -> List[T]:
"""Execute a function on all workers synchronously.
Args:
func: The function to execute on each worker.
*args: Positional arguments to pass to the function.
**kwargs: Keyword arguments to pass to the function.
Returns:
A list of results from each worker, in worker rank order.
"""
pass
@abc.abstractmethod
def execute_async(self, func: Callable[..., T], *args, **kwargs) -> List[ObjectRef]:
"""Execute a function on all workers asynchronously.
Args:
func: The function to execute on each worker.
*args: Positional arguments to pass to the function.
**kwargs: Keyword arguments to pass to the function.
Returns:
A list of ObjectRef results from each worker, in worker rank order.
"""
pass
@abc.abstractmethod
def execute_single(
self, worker_index: int, func: Callable[..., T], *args, **kwargs
) -> T:
"""Execute a function on a single worker synchronously.
Args:
worker_index: The index of the worker to execute on.
func: The function to execute.
*args: Positional arguments to pass to the function.
**kwargs: Keyword arguments to pass to the function.
Returns:
The result from the specified worker.
"""
pass
@abc.abstractmethod
def execute_single_async(
self, worker_index: int, func: Callable[..., T], *args, **kwargs
) -> ObjectRef:
"""Execute a function on a single worker asynchronously.
Args:
worker_index: The index of the worker to execute on.
func: The function to execute.
*args: Positional arguments to pass to the function.
**kwargs: Keyword arguments to pass to the function.
Returns:
An ObjectRef to the result from the specified worker.
"""
pass
@abc.abstractmethod
def __len__(self) -> int:
"""Return the number of workers in the group."""
pass
@abc.abstractmethod
def get_resources_per_worker(self) -> dict:
"""Get the resources allocated per worker.
Returns:
A dictionary mapping resource names to quantities per worker.
Common keys include "CPU", "GPU", "memory".
"""
pass
python/ray/train/_internal/checkpoint_manager.py
+206
View File
@@ -0,0 +1,206 @@
import logging
import numbers
from typing import Any, Callable, Dict, List, Optional, Tuple
from ray._private import ray_constants
from ray._private.dict import flatten_dict
from ray.air._internal.util import is_nan
from ray.air.config import MAX
from ray.train import Checkpoint, CheckpointConfig
from ray.train._internal.session import _TrainingResult
from ray.train._internal.storage import _delete_fs_path
from ray.train.constants import TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE
logger = logging.getLogger(__name__)
def _insert_into_sorted_list(
list: List[_TrainingResult],
item: _TrainingResult,
key: Callable[[_TrainingResult], Any],
checkpoint_to_report_index: Optional[Dict[Checkpoint, int]] = None,
):
"""Insert an item into a sorted list with a custom key function.
Args:
list: The list to insert the item into.
item: The item to insert.
key: The key function to use to sort the list.
checkpoint_to_report_index: A dictionary mapping checkpoints to report indices.
Used to break ties when scores are equal.
"""
checkpoint_to_report_index = checkpoint_to_report_index or {}
# TODO: optimize this with sortedlist, batching, etc
i = 0
while i < len(list):
# When scores are equal, later checkpoints are later in the list.
list_item_key, item_key = key(list[i]), key(item)
if list_item_key > item_key or (
list_item_key == item_key
and checkpoint_to_report_index.get(list[i].checkpoint, 0)
> checkpoint_to_report_index.get(item.checkpoint, 0)
):
break
i += 1
list.insert(i, item)
class _CheckpointManager:
"""Checkpoint manager that handles checkpoint book-keeping for a trial.
The main purpose of this abstraction is to keep the top K checkpoints based on
recency/a user-provided metric.
NOTE: This class interacts with `_TrainingResult` objects, which are
(checkpoint, metrics) pairs. This is to order checkpoints by metrics.
Args:
checkpoint_config: Defines how many and which checkpoints to keep.
"""
def __init__(self, checkpoint_config: Optional[CheckpointConfig]):
self._checkpoint_config = checkpoint_config or CheckpointConfig()
# List of checkpoints ordered by ascending score.
self._checkpoint_results: List[_TrainingResult] = []
# The latest registered checkpoint.
# This should never be immediately deleted upon registration,
# even if it's not in the top K checkpoints, based on score.
self._latest_checkpoint_result: Optional[_TrainingResult] = None
if (
self._checkpoint_config.num_to_keep is not None
and self._checkpoint_config.num_to_keep <= 0
):
raise ValueError(
f"`num_to_keep` must >= 1, got: "
f"{self._checkpoint_config.num_to_keep}"
)
@property
def checkpoint_config(self):
return self._checkpoint_config
def register_checkpoint(self, checkpoint_result: _TrainingResult):
"""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:
checkpoint_result: Tracked training result containing the checkpoint
and associated metrics to add to bookkeeping.
"""
self._latest_checkpoint_result = checkpoint_result
score_attr = self._checkpoint_config.checkpoint_score_attribute
if ray_constants.env_bool(TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE, False):
metrics = (
{score_attr: checkpoint_result.metrics[score_attr]}
if score_attr in checkpoint_result.metrics
else {}
)
checkpoint_result = _TrainingResult(
checkpoint=checkpoint_result.checkpoint,
metrics=metrics,
)
if score_attr is not None and score_attr in checkpoint_result.metrics:
# 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,
)
else:
# If no metric is provided, just append (ordering by time of registration).
self._checkpoint_results.append(checkpoint_result)
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.
results_to_delete = worst_results - {self._latest_checkpoint_result}
# 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:
checkpoint = checkpoint_result.checkpoint
logger.debug("Deleting checkpoint: %s", checkpoint)
_delete_fs_path(fs=checkpoint.filesystem, fs_path=checkpoint.path)
def _get_checkpoint_score(
self, checkpoint: _TrainingResult
) -> Tuple[bool, numbers.Number]:
"""Get the score for a checkpoint, according to checkpoint config.
If `mode="min"`, the metric is negated so that the lowest score is
treated as the best.
Args:
checkpoint: The training result whose metrics should be scored.
Returns:
Tuple: A tuple of (not_is_nan: bool, score: numbers.Number).
This score orders: nan values < float("-inf") < valid numeric metrics
"""
checkpoint_score_attribute = self._checkpoint_config.checkpoint_score_attribute
if checkpoint_score_attribute:
flat_metrics = flatten_dict(checkpoint.metrics)
try:
checkpoint_result = flat_metrics[checkpoint_score_attribute]
except KeyError:
valid_keys = list(flat_metrics.keys())
logger.error(
f"Result dict has no key: {checkpoint_score_attribute}. "
f"checkpoint_score_attr must be set to a key in the "
f"result dict. Valid keys are: {valid_keys}"
)
checkpoint_result = float("-inf")
else:
checkpoint_result = float("-inf")
checkpoint_score_order = self._checkpoint_config.checkpoint_score_order
order_factor = 1.0 if checkpoint_score_order == MAX else -1.0
checkpoint_score = order_factor * checkpoint_result
if not isinstance(checkpoint_score, numbers.Number):
raise ValueError(
f"Unable to persist checkpoint for "
f"checkpoint_score_attribute: "
f"{checkpoint_score_attribute} with value "
f"{checkpoint_score}. "
f"This attribute must be numerical."
)
return (
(not is_nan(checkpoint_score), checkpoint_score)
if not is_nan(checkpoint_score)
else (False, float("-inf"))
)
@property
def best_checkpoint_result(self) -> Optional[_TrainingResult]:
return self._checkpoint_results[-1] if self._checkpoint_results else None
@property
def latest_checkpoint_result(self) -> Optional[_TrainingResult]:
return self._latest_checkpoint_result
@property
def best_checkpoint_results(self) -> List[_TrainingResult]:
if self._checkpoint_config.num_to_keep is None:
return self._checkpoint_results
return self._checkpoint_results[-self._checkpoint_config.num_to_keep :]
python/ray/train/_internal/data_config.py
+178
View File
@@ -0,0 +1,178 @@
import copy
from collections import defaultdict
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
from ray.actor import ActorHandle
from ray.util.annotations import DeveloperAPI, PublicAPI
if TYPE_CHECKING:
from ray.data import DataIterator, Dataset, ExecutionOptions, NodeIdStr
@PublicAPI(stability="stable")
class DataConfig:
"""Class responsible for configuring Train dataset preprocessing.
For advanced use cases, this class can be subclassed and the `configure()` method
overridden for custom data preprocessing.
"""
def __init__(
self,
datasets_to_split: Union[Literal["all"], List[str]] = "all",
execution_options: Optional[
Union["ExecutionOptions", Dict[str, "ExecutionOptions"]]
] = None,
enable_shard_locality: bool = True,
):
"""Construct a DataConfig.
Args:
datasets_to_split: Specifies which datasets should be split among workers.
Can be set to "all" or a list of dataset names. Defaults to "all",
i.e. split all datasets.
execution_options: The execution options to pass to Ray Data. Can be either:
1. A single ExecutionOptions object that is applied to all datasets.
2. A dict mapping dataset names to ExecutionOptions. If a dataset name
is not in the dict, it defaults to ``DataConfig.default_ingest_options()``.
By default, the options are optimized for data ingest. When overriding,
base your options off ``DataConfig.default_ingest_options()``.
enable_shard_locality: If true, dataset sharding across Train workers will
consider locality to minimize cross-node data transfer. Enabled by default.
"""
from ray.data import ExecutionOptions
if isinstance(datasets_to_split, list) or datasets_to_split == "all":
self._datasets_to_split = datasets_to_split
else:
raise TypeError(
"`datasets_to_split` should be a 'all' or a list of strings of "
"dataset names. Received "
f"{type(datasets_to_split).__name__} with value {datasets_to_split}."
)
default_execution_options = DataConfig.default_ingest_options()
if isinstance(execution_options, ExecutionOptions):
default_execution_options = execution_options
# If None, all datasets will use the default ingest options.
self._execution_options: Dict[str, "ExecutionOptions"] = defaultdict(
lambda: copy.deepcopy(default_execution_options)
)
if isinstance(execution_options, dict):
self._execution_options.update(execution_options)
self._enable_shard_locality = enable_shard_locality
self._num_train_cpus = 0.0
self._num_train_gpus = 0.0
def set_train_total_resources(self, num_train_cpus: float, num_train_gpus: float):
"""Set the total number of CPUs and GPUs used by training.
If CPU or GPU resource limits are not set, they will be set to the
total cluster resources minus the resources used by training.
"""
# TODO: We may also include other resources besides CPU and GPU.
self._num_train_cpus = num_train_cpus
self._num_train_gpus = num_train_gpus
def _get_execution_options(self, dataset_name: str) -> "ExecutionOptions":
"""Return a copy of the configured execution options for a given dataset name."""
return copy.deepcopy(self._execution_options[dataset_name])
@DeveloperAPI
def configure(
self,
datasets: Dict[str, "Dataset"],
world_size: int,
worker_handles: Optional[List[ActorHandle]],
worker_node_ids: Optional[List["NodeIdStr"]],
**kwargs,
) -> List[Dict[str, "DataIterator"]]:
"""Configure how Train datasets should be assigned to workers.
Args:
datasets: The datasets dict passed to Train by the user.
world_size: The number of Train workers in total.
worker_handles: The actor handles of the Train workers.
worker_node_ids: The node ids of the Train workers.
**kwargs: Forwards compatibility placeholder.
Returns:
A list of dataset splits for each worker. The size of the list must be
equal to `world_size`. Each element of the list contains the assigned
`DataIterator` instances by name for the worker.
"""
from ray.data._internal.execution.interfaces.execution_options import (
ExecutionResources,
)
output = [{} for _ in range(world_size)]
for dataset_name, dataset in datasets.items():
if dataset.name is None:
dataset.set_name(dataset_name)
if self._datasets_to_split == "all":
datasets_to_split = set(datasets.keys())
else:
datasets_to_split = set(self._datasets_to_split)
locality_hints = worker_node_ids if self._enable_shard_locality else None
for name, ds in datasets.items():
execution_options = self._get_execution_options(name)
if execution_options.is_resource_limits_default():
if not self._scaling_policy_reserves_train_resources():
execution_options.exclude_resources = (
execution_options.exclude_resources.add(
ExecutionResources(
cpu=self._num_train_cpus, gpu=self._num_train_gpus
)
)
)
ds = ds.copy(ds)
ds.context.execution_options = execution_options
if name in datasets_to_split:
for i, split in enumerate(
ds.streaming_split(
world_size, equal=True, locality_hints=locality_hints
)
):
output[i][name] = split
else:
for i in range(world_size):
output[i][name] = ds.iterator()
return output
@classmethod
def _scaling_policy_reserves_train_resources(cls) -> bool:
"""True iff Ray Train V2's ScalingPolicy will register training resources
with the AutoscalingCoordinator for this run.
"""
from ray.train.v2._internal.constants import is_v2_enabled
return is_v2_enabled()
@staticmethod
def default_ingest_options() -> "ExecutionOptions":
"""The default Ray Data options used for data ingest.
By default, configurations are carried over from what is already set
in DataContext.
"""
from ray.data import ExecutionOptions
from ray.data.context import DataContext
ctx = DataContext.get_current()
return ExecutionOptions(
resource_limits=ctx.execution_options.resource_limits,
exclude_resources=ctx.execution_options.exclude_resources,
preserve_order=ctx.execution_options.preserve_order,
verbose_progress=ctx.execution_options.verbose_progress,
)
python/ray/train/_internal/dl_predictor.py
View File
python/ray/train/_internal/framework_checkpoint.py
+57
View File
@@ -0,0 +1,57 @@
from typing import TYPE_CHECKING, Optional
import ray.cloudpickle as ray_pickle
from ray._common.utils import binary_to_hex, hex_to_binary
from ray.train._checkpoint import Checkpoint
if TYPE_CHECKING:
from ray.data.preprocessor import Preprocessor
PREPROCESSOR_KEY = "preprocessor_pkl"
class FrameworkCheckpoint(Checkpoint):
"""A checkpoint to preserve the functionality of legacy
framework-specific checkpoints.
Example:
>>> import tempfile
>>> from ray.data.preprocessor import Preprocessor
>>> checkpoint = FrameworkCheckpoint(tempfile.mkdtemp())
>>> checkpoint.get_preprocessor() is None
True
>>> preprocessor = Preprocessor()
>>> preprocessor._attr = 1234
>>> checkpoint.set_preprocessor(preprocessor)
>>> checkpoint.get_preprocessor()._attr
1234
"""
def get_preprocessor(self) -> Optional["Preprocessor"]:
"""Return the preprocessor stored in the checkpoint.
.. warning::
The checkpoint path must point to a **trusted** source.
The preprocessor is stored as a pickle blob inside the checkpoint
metadata. Loading a checkpoint from an untrusted path (shared
storage, downloaded artifact, checkpoint produced by a different
party) is equivalent to executing arbitrary Python code. Never
call this method on a checkpoint you do not fully control.
Returns:
The preprocessor stored in the checkpoint, or ``None`` if no
preprocessor was stored.
"""
metadata = self.get_metadata()
preprocessor_bytes = metadata.get(PREPROCESSOR_KEY)
if preprocessor_bytes is None:
return None
return ray_pickle.loads(hex_to_binary(preprocessor_bytes))
def set_preprocessor(self, preprocessor: "Preprocessor"):
"""Store a preprocessor with the checkpoint."""
self.update_metadata(
{PREPROCESSOR_KEY: binary_to_hex(ray_pickle.dumps(preprocessor))}
)
python/ray/train/_internal/session.py
+1219
View File
File diff suppressed because it is too large Load Diff
python/ray/train/_internal/state/__init__.py
+14
View File
@@ -0,0 +1,14 @@
from ray.train._internal.state.state_manager import TrainRunStateManager
try:
import pydantic # noqa: F401
except ImportError:
raise ModuleNotFoundError(
"pydantic isn't installed."
"To install pydantic, please run 'pip install pydantic'"
)
__all__ = [
"TrainRunStateManager",
]
python/ray/train/_internal/state/export.py
+105
View File
@@ -0,0 +1,105 @@
from typing import Optional
from ray.core.generated.export_train_state_pb2 import (
ExportTrainRunAttemptEventData as ProtoTrainRunAttempt,
ExportTrainRunEventData as ProtoTrainRun,
)
from ray.train._internal.state.schema import (
ActorStatusEnum,
RunStatusEnum,
TrainRunInfo,
TrainWorkerInfo,
)
TRAIN_SCHEMA_VERSION = 1
RAY_TRAIN_VERSION = 1
# Status mapping dictionaries
_ACTOR_STATUS_MAP = {
ActorStatusEnum.ALIVE: ProtoTrainRunAttempt.ActorStatus.ALIVE,
ActorStatusEnum.DEAD: ProtoTrainRunAttempt.ActorStatus.DEAD,
}
_RUN_ATTEMPT_STATUS_MAP = {
RunStatusEnum.STARTED: ProtoTrainRunAttempt.RunAttemptStatus.PENDING,
RunStatusEnum.RUNNING: ProtoTrainRunAttempt.RunAttemptStatus.RUNNING,
RunStatusEnum.FINISHED: ProtoTrainRunAttempt.RunAttemptStatus.FINISHED,
RunStatusEnum.ERRORED: ProtoTrainRunAttempt.RunAttemptStatus.ERRORED,
RunStatusEnum.ABORTED: ProtoTrainRunAttempt.RunAttemptStatus.ABORTED,
}
_RUN_STATUS_MAP = {
RunStatusEnum.STARTED: ProtoTrainRun.RunStatus.INITIALIZING,
RunStatusEnum.RUNNING: ProtoTrainRun.RunStatus.RUNNING,
RunStatusEnum.FINISHED: ProtoTrainRun.RunStatus.FINISHED,
RunStatusEnum.ERRORED: ProtoTrainRun.RunStatus.ERRORED,
RunStatusEnum.ABORTED: ProtoTrainRun.RunStatus.ABORTED,
}
def _ms_to_ns(ms: Optional[int]) -> Optional[int]:
if ms is None:
return None
return ms * 1000000
# Helper conversion functions
def _to_proto_resources(resources: dict) -> ProtoTrainRunAttempt.TrainResources:
"""Convert resources dictionary to protobuf TrainResources."""
return ProtoTrainRunAttempt.TrainResources(resources=resources)
def _to_proto_worker(worker: TrainWorkerInfo) -> ProtoTrainRunAttempt.TrainWorker:
"""Convert TrainWorker to protobuf format."""
proto_worker = ProtoTrainRunAttempt.TrainWorker(
world_rank=worker.world_rank,
local_rank=worker.local_rank,
node_rank=worker.node_rank,
actor_id=bytes.fromhex(worker.actor_id),
node_id=bytes.fromhex(worker.node_id),
node_ip=worker.node_ip,
pid=worker.pid,
gpu_ids=worker.gpu_ids,
status=_ACTOR_STATUS_MAP[worker.status],
resources=_to_proto_resources(worker.resources),
)
return proto_worker
# Main conversion functions
def train_run_info_to_proto_run(run_info: TrainRunInfo) -> ProtoTrainRun:
"""Convert TrainRunInfo to TrainRun protobuf format."""
proto_run = ProtoTrainRun(
schema_version=TRAIN_SCHEMA_VERSION,
ray_train_version=RAY_TRAIN_VERSION,
id=run_info.id,
name=run_info.name,
job_id=bytes.fromhex(run_info.job_id),
controller_actor_id=bytes.fromhex(run_info.controller_actor_id),
status=_RUN_STATUS_MAP[run_info.run_status],
status_detail=run_info.status_detail,
start_time_ns=_ms_to_ns(run_info.start_time_ms),
end_time_ns=_ms_to_ns(run_info.end_time_ms),
)
return proto_run
def train_run_info_to_proto_attempt(run_info: TrainRunInfo) -> ProtoTrainRunAttempt:
"""Convert TrainRunInfo to TrainRunAttempt protobuf format."""
proto_attempt = ProtoTrainRunAttempt(
schema_version=TRAIN_SCHEMA_VERSION,
ray_train_version=RAY_TRAIN_VERSION,
run_id=run_info.id,
attempt_id=run_info.id, # Same as run_id
status=_RUN_ATTEMPT_STATUS_MAP[run_info.run_status],
status_detail=run_info.status_detail,
start_time_ns=_ms_to_ns(run_info.start_time_ms),
end_time_ns=_ms_to_ns(run_info.end_time_ms),
resources=[_to_proto_resources(r) for r in run_info.resources],
workers=[_to_proto_worker(worker) for worker in run_info.workers],
)
return proto_attempt
python/ray/train/_internal/state/schema.py
+165
View File
@@ -0,0 +1,165 @@
from enum import Enum
from typing import Dict, List, Optional
from ray._common.pydantic_compat import BaseModel, Field
from ray.dashboard.modules.job.pydantic_models import JobDetails
from ray.util.annotations import DeveloperAPI
MAX_ERROR_STACK_TRACE_LENGTH = 50000
@DeveloperAPI
class RunStatusEnum(str, Enum):
"""Enumeration for the status of a train run."""
# (Deprecated) Replaced by RUNNING.
# The train run has started
STARTED = "STARTED"
# The train run is running
RUNNING = "RUNNING"
# The train run was terminated as expected
FINISHED = "FINISHED"
# The train run was terminated early due to errors in the training function
ERRORED = "ERRORED"
# The train run was terminated early due to system errors or controller errors
ABORTED = "ABORTED"
@DeveloperAPI
class ActorStatusEnum(str, Enum):
DEAD = "DEAD"
ALIVE = "ALIVE"
@DeveloperAPI
class TrainWorkerInfo(BaseModel):
"""Metadata of a Ray Train worker."""
actor_id: str = Field(description="Actor ID of the worker.")
world_rank: int = Field(description="World rank of the worker.")
local_rank: int = Field(description="Local rank of the worker.")
node_rank: int = Field(description="Node rank of the worker.")
node_id: str = Field(description="ID of the node that the worker is running on.")
node_ip: str = Field(
description="IP address of the node that the worker is running on."
)
pid: int = Field(description="Process ID of the worker.")
gpu_ids: List[int] = Field(
description="A list of GPU ids allocated to that worker."
)
status: ActorStatusEnum = Field(
description="The status of the train worker actor. It can be ALIVE or DEAD."
)
resources: Dict[str, float] = Field(
description="The resources allocated to the worker."
)
@DeveloperAPI
class MemoryInfo(BaseModel):
rss: int
vms: int
pfaults: Optional[int] = None
pageins: Optional[int] = None
@DeveloperAPI
class ProcessStats(BaseModel):
cpuPercent: float
# total memory, free memory, memory used ratio
mem: Optional[List[int]] = None
memoryInfo: MemoryInfo
class ProcessGPUUsage(BaseModel):
# This gpu usage stats from a process
pid: int
gpuMemoryUsage: int
@DeveloperAPI
class GPUStats(BaseModel):
uuid: str
index: int
name: str
utilizationGpu: Optional[float] = None
memoryUsed: float
memoryTotal: float
processInfo: ProcessGPUUsage
@DeveloperAPI
class TrainWorkerInfoWithDetails(TrainWorkerInfo):
"""Metadata of a Ray Train worker."""
processStats: Optional[ProcessStats] = Field(
None, description="Process stats of the worker."
)
gpus: List[GPUStats] = Field(
default_factory=list,
description=(
"GPU stats of the worker. "
"Only returns GPUs that are attached to the worker process."
),
)
@DeveloperAPI
class TrainDatasetInfo(BaseModel):
name: str = Field(
description="The key of the dataset dict specified in Ray Train Trainer."
)
dataset_uuid: str = Field(description="The uuid of the dataset.")
dataset_name: Optional[str] = Field(None, description="The name of the dataset.")
@DeveloperAPI
class TrainRunInfo(BaseModel):
"""Metadata for a Ray Train run and information about its workers."""
name: str = Field(description="The name of the Train run.")
id: str = Field(description="The unique identifier for each Train run.")
job_id: str = Field(description="The Ray Job ID.")
controller_actor_id: str = Field(description="Actor Id of the Train controller.")
workers: List[TrainWorkerInfo] = Field(
description="A List of Train workers sorted by global ranks."
)
datasets: List[TrainDatasetInfo] = Field(
description="A List of dataset info for this Train run."
)
run_status: RunStatusEnum = Field(
description="The current status of the train run. It can be one of the "
"following: RUNNING, FINISHED, ERRORED, or ABORTED."
)
status_detail: str = Field(
description="Detailed information about the current run status, "
"such as error messages."
)
start_time_ms: int = Field(
description="The UNIX timestamp of the start time of this Train run."
)
end_time_ms: Optional[int] = Field(
None,
description="The UNIX timestamp of the end time of this Train run. "
"If null, the Train run has not ended yet.",
)
resources: List[Dict[str, float]] = Field(
description="The resources allocated to the worker."
)
@DeveloperAPI
class TrainRunInfoWithDetails(TrainRunInfo):
"""Metadata for a Ray Train run and information about its workers."""
workers: List[TrainWorkerInfoWithDetails] = Field(
description="A List of Train workers sorted by global ranks."
)
job_details: Optional[JobDetails] = Field(
None, description="Details of the job that started this Train run."
)
@DeveloperAPI
class TrainRunsResponse(BaseModel):
train_runs: List[TrainRunInfoWithDetails]
python/ray/train/_internal/state/state_actor.py
+152
View File
@@ -0,0 +1,152 @@
import logging
import os
import threading
from typing import Dict, Optional
import ray
from ray._private.event.export_event_logger import (
EventLogType,
check_export_api_enabled,
get_export_event_logger,
)
from ray.actor import ActorHandle
from ray.train._internal.state.schema import TrainRunInfo
logger = logging.getLogger(__name__)
@ray.remote(num_cpus=0)
class TrainStateActor:
def __init__(self):
self._run_infos: Dict[str, TrainRunInfo] = {}
(
self._export_logger,
self._is_train_run_export_api_enabled,
self._is_train_run_attempt_export_api_enabled,
) = self._init_export_logger()
def register_train_run(self, run_info: TrainRunInfo) -> None:
# Register a new train run.
self._run_infos[run_info.id] = run_info
self._maybe_export_train_run(run_info)
self._maybe_export_train_run_attempt(run_info)
def get_train_run(self, run_id: str) -> Optional[TrainRunInfo]:
# Retrieve a registered run with its id
return self._run_infos.get(run_id, None)
def get_all_train_runs(self) -> Dict[str, TrainRunInfo]:
# Retrieve all registered train runs
return self._run_infos
# ============================
# Export API
# ============================
def is_export_api_enabled(self) -> bool:
return self._export_logger is not None
def _init_export_logger(self) -> tuple[Optional[logging.Logger], bool, bool]:
"""Initialize the export logger and check if the export API is enabled.
Returns:
A tuple containing:
- The export logger (or None if export API is not enabled).
- A boolean indicating if the export API is enabled for train runs.
- A boolean indicating if the export API is enabled for train run attempts.
"""
# Proto schemas should be imported within the scope of TrainStateActor to
# prevent serialization errors.
from ray.core.generated.export_event_pb2 import ExportEvent
is_train_run_export_api_enabled = check_export_api_enabled(
ExportEvent.SourceType.EXPORT_TRAIN_RUN
)
is_train_run_attempt_export_api_enabled = check_export_api_enabled(
ExportEvent.SourceType.EXPORT_TRAIN_RUN_ATTEMPT
)
export_api_enabled = (
is_train_run_export_api_enabled or is_train_run_attempt_export_api_enabled
)
if not export_api_enabled:
return None, False, False
log_directory = os.path.join(
ray._private.worker._global_node.get_session_dir_path(), "logs"
)
logger = None
try:
logger = get_export_event_logger(
EventLogType.TRAIN_STATE,
log_directory,
)
except Exception:
logger.exception(
"Unable to initialize the export event logger, so no Train export "
"events will be written."
)
if logger is None:
return None, False, False
return (
logger,
is_train_run_export_api_enabled,
is_train_run_attempt_export_api_enabled,
)
def _maybe_export_train_run(self, run_info: TrainRunInfo) -> None:
if not self._is_train_run_export_api_enabled:
return
from ray.train._internal.state.export import train_run_info_to_proto_run
run_proto = train_run_info_to_proto_run(run_info)
self._export_logger.send_event(run_proto)
def _maybe_export_train_run_attempt(self, run_info: TrainRunInfo) -> None:
if not self._is_train_run_attempt_export_api_enabled:
return
from ray.train._internal.state.export import train_run_info_to_proto_attempt
run_attempt_proto = train_run_info_to_proto_attempt(run_info)
self._export_logger.send_event(run_attempt_proto)
TRAIN_STATE_ACTOR_NAME = "train_state_actor"
TRAIN_STATE_ACTOR_NAMESPACE = "_train_state_actor"
_state_actor_lock: threading.RLock = threading.RLock()
def get_or_create_state_actor() -> ActorHandle:
"""Get or create a `TrainStateActor` on the head node."""
with _state_actor_lock:
state_actor = TrainStateActor.options(
name=TRAIN_STATE_ACTOR_NAME,
namespace=TRAIN_STATE_ACTOR_NAMESPACE,
get_if_exists=True,
lifetime="detached",
resources={"node:__internal_head__": 0.001},
# Escape from the parent's placement group
scheduling_strategy="DEFAULT",
).remote()
# Ensure the state actor is ready
ray.get(state_actor.__ray_ready__.remote())
return state_actor
def get_state_actor() -> Optional[ActorHandle]:
"""Get the `TrainStateActor` if exists, otherwise return None."""
try:
return ray.get_actor(
name=TRAIN_STATE_ACTOR_NAME,
namespace=TRAIN_STATE_ACTOR_NAMESPACE,
)
except ValueError:
return None
python/ray/train/_internal/state/state_manager.py
+132
View File
@@ -0,0 +1,132 @@
import logging
import os
from collections import defaultdict
from typing import TYPE_CHECKING, Any, Dict, List
import ray
from ray.train._internal.state.schema import (
ActorStatusEnum,
RunStatusEnum,
TrainDatasetInfo,
TrainRunInfo,
TrainWorkerInfo,
)
from ray.train._internal.utils import check_for_failure
from ray.train._internal.worker_group import WorkerGroup
if TYPE_CHECKING:
from ray.data import Dataset
logger = logging.getLogger(__name__)
class TrainRunStateManager:
"""A class that aggregates and reports train run info to TrainStateActor.
This manager class is created on the train controller layer for each run.
"""
def __init__(self, state_actor) -> None:
self.state_actor = state_actor
self.train_run_info_dict = defaultdict(dict)
def register_train_run(
self,
run_id: str,
job_id: str,
run_name: str,
run_status: str,
controller_actor_id: str,
datasets: Dict[str, "Dataset"],
worker_group: WorkerGroup,
start_time_ms: float,
resources: List[Dict[str, float]],
status_detail: str = "",
) -> None:
"""Collect Train Run Info and report to StateActor."""
if not self.state_actor:
logger.warning(
"Unable to register train run since `TrainStateActor` is not started."
)
return
def collect_train_worker_info():
train_context = ray.train.get_context()
core_context = ray.runtime_context.get_runtime_context()
return TrainWorkerInfo(
world_rank=train_context.get_world_rank(),
local_rank=train_context.get_local_rank(),
node_rank=train_context.get_node_rank(),
actor_id=core_context.get_actor_id(),
node_id=core_context.get_node_id(),
node_ip=ray.util.get_node_ip_address(),
gpu_ids=ray.get_gpu_ids(),
pid=os.getpid(),
resources=resources[0],
status=ActorStatusEnum.ALIVE,
)
futures = [
worker_group.execute_single_async(index, collect_train_worker_info)
for index in range(len(worker_group))
]
success, exception = check_for_failure(futures)
if not success:
logger.error(
"Failed to collect run information from the Ray Train "
f"workers:\n{exception}"
)
return
worker_info_list = ray.get(futures)
worker_info_list = sorted(worker_info_list, key=lambda info: info.world_rank)
dataset_info_list = [
TrainDatasetInfo(
name=ds_name,
dataset_name=ds._dataset_name,
dataset_uuid=ds._uuid,
)
for ds_name, ds in datasets.items()
]
updates = dict(
id=run_id,
job_id=job_id,
name=run_name,
controller_actor_id=controller_actor_id,
workers=worker_info_list,
datasets=dataset_info_list,
start_time_ms=start_time_ms,
run_status=run_status,
status_detail=status_detail,
resources=resources,
)
# Clear the cached info to avoid registering the same run twice
self.train_run_info_dict[run_id] = {}
self._update_train_run_info(run_id, updates)
def end_train_run(
self,
run_id: str,
run_status: RunStatusEnum,
status_detail: str,
end_time_ms: int,
):
"""Update the train run status when the training is finished."""
updates = dict(
run_status=run_status,
status_detail=status_detail,
end_time_ms=end_time_ms,
)
self._update_train_run_info(run_id, updates)
def _update_train_run_info(self, run_id: str, updates: Dict[str, Any]) -> None:
"""Update specific fields of a registered TrainRunInfo instance."""
if run_id in self.train_run_info_dict:
self.train_run_info_dict[run_id].update(updates)
train_run_info = TrainRunInfo(**self.train_run_info_dict[run_id])
ray.get(self.state_actor.register_train_run.remote(train_run_info))
python/ray/train/_internal/storage.py
+736
View File
@@ -0,0 +1,736 @@
# 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
try:
# check if Arrow has S3 support
from pyarrow.fs import S3FileSystem
except ImportError:
S3FileSystem = None
# isort: on
import fnmatch
import logging
import os
import shutil
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Type, Union
from ray.air._internal.filelock import TempFileLock
from ray.train._internal.syncer import SyncConfig, Syncer, _BackgroundSyncer
from ray.train.constants import _get_ray_train_session_dir
from ray.util.annotations import DeveloperAPI
if TYPE_CHECKING:
from ray.train._checkpoint import Checkpoint
logger = logging.getLogger(__name__)
_VALIDATE_STORAGE_MARKER_FILENAME = ".validate_storage_marker"
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: Forwarded to the ``fsspec.implementations.local.LocalFileSystem``
parent class.
"""
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 S3FileSystem and 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):
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: Optional[Callable[[pyarrow.fs.FileInfo], bool]] = None,
) -> 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.
"""
if file_filter is None:
file_filter = lambda x: True # noqa: E731
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 use.
fs_path: The path on the filesystem to check.
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 tuple of (filesystem, path) resolved from the inputs.
"""
storage_path = str(storage_path)
if storage_filesystem:
return storage_filesystem, storage_path
return pyarrow.fs.FileSystem.from_uri(storage_path)
class _FilesystemSyncer(_BackgroundSyncer):
"""Syncer between local filesystem and a `storage_filesystem`."""
def __init__(self, storage_filesystem: Optional["pyarrow.fs.FileSystem"], **kwargs):
self.storage_filesystem = storage_filesystem
super().__init__(**kwargs)
def _sync_up_command(
self, local_path: str, uri: str, exclude: Optional[List] = None
) -> Tuple[Callable, Dict]:
# TODO(justinvyu): Defer this cleanup up as part of the
# external-facing Syncer deprecation.
fs_path = uri
return (
_upload_to_fs_path,
dict(
local_path=local_path,
fs=self.storage_filesystem,
fs_path=fs_path,
exclude=exclude,
),
)
def _sync_down_command(self, uri: str, local_path: str) -> Tuple[Callable, Dict]:
fs_path = uri
return (
_download_from_fs_path,
dict(
fs=self.storage_filesystem,
fs_path=fs_path,
local_path=local_path,
),
)
def _delete_command(self, uri: str) -> Tuple[Callable, Dict]:
fs_path = uri
return _delete_fs_path, dict(fs=self.storage_filesystem, fs_path=fs_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'
>>> storage.current_checkpoint_index = 1
>>> storage.checkpoint_fs_path
'bucket/path/exp_name/trial_dir/checkpoint_000001'
>>> 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,
sync_config: Optional[SyncConfig] = None,
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
trial_dir_name: Optional[str] = None,
current_checkpoint_index: int = -1,
):
from ray.tune.utils import date_str
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.trial_dir_name = trial_dir_name
self.current_checkpoint_index = current_checkpoint_index
self.sync_config = sync_config or SyncConfig()
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.syncer: Syncer = _FilesystemSyncer(
storage_filesystem=self.storage_filesystem,
sync_period=self.sync_config.sync_period,
sync_timeout=self.sync_config.sync_timeout,
)
self._create_validation_file()
self._check_validation_file()
# Timestamp is used to create a unique session directory for the current
# training job. This is used to avoid conflicts when multiple training jobs
# run with the same name in the same cluster.
# This is set ONCE at the creation of the storage context, on the driver.
self._timestamp = date_str()
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"
f" trial_dir_name='{self.trial_dir_name}',\n"
f" current_checkpoint_index={self.current_checkpoint_index},\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 _update_checkpoint_index(self, metrics: Dict):
# Per default, increase by 1. This can be overwritten to customize checkpoint
# directories.
self.current_checkpoint_index += 1
def persist_current_checkpoint(self, checkpoint: "Checkpoint") -> "Checkpoint":
"""Persists a given checkpoint to the current checkpoint path on the filesystem.
"Current" is defined by the `current_checkpoint_index` attribute of the
storage context.
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, checkpoint_fs_path).
Returns:
Checkpoint: A Checkpoint pointing to the persisted checkpoint location.
"""
# TODO(justinvyu): Fix this cyclical import.
logger.debug(
"Copying checkpoint files to storage path:\n"
"({source_fs}, {source}) -> ({dest_fs}, {destination})".format(
source=checkpoint.path,
destination=self.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(self.checkpoint_fs_path)
_pyarrow_fs_copy_files(
source=checkpoint.path,
destination=self.checkpoint_fs_path,
source_filesystem=checkpoint.filesystem,
destination_filesystem=self.storage_filesystem,
)
persisted_checkpoint = checkpoint.__class__(
filesystem=self.storage_filesystem,
path=self.checkpoint_fs_path,
)
logger.info(f"Checkpoint successfully created at: {persisted_checkpoint}")
return persisted_checkpoint
def persist_artifacts(self, force: bool = False) -> None:
"""Persists all artifacts within `trial_local_dir` to storage.
This method possibly launches a background task to sync the trial dir,
depending on the `sync_period` + `sync_artifacts_on_checkpoint`
settings of `SyncConfig`.
`(local_fs, trial_working_dir) -> (storage_filesystem, trial_fs_path)`
Args:
force: If True, wait for a previous sync to finish, launch a new one,
and wait for that one to finish. By the end of a `force=True` call, the
latest version of the trial artifacts will be persisted.
"""
if not self.sync_config.sync_artifacts:
return
# Skip if there are no artifacts to sync
is_empty = not any(os.scandir(self.trial_working_directory))
if is_empty:
return
if force:
self.syncer.wait()
self.syncer.sync_up(
local_dir=self.trial_working_directory, remote_dir=self.trial_fs_path
)
self.syncer.wait()
else:
self.syncer.sync_up_if_needed(
local_dir=self.trial_working_directory, remote_dir=self.trial_fs_path
)
@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()
def _get_session_path(self) -> str:
"""The Ray Train/Tune session local directory used to stage files
before persisting to the storage filesystem."""
return Path(
_get_ray_train_session_dir(), self._timestamp, self.experiment_dir_name
).as_posix()
@property
def experiment_driver_staging_path(self) -> str:
"""The local filesystem path of the experiment directory on the driver node.
The driver is the node where `Trainer.fit`/`Tuner.fit` is being called.
This path is of the form:
`/tmp/ray/session_<session_id>/artifacts/<ray-train-job-timestamp>/
<experiment_dir_name>/driver_artifacts`
This should be used as the temporary staging location for files *on the driver*
before syncing them to `experiment_fs_path`.
For example, the search algorithm should dump its state to this directory.
See `trial_driver_staging_path` for writing trial-specific artifacts.
The directory is synced to
`{storage_path}/{experiment_dir_name}` periodically.
See `_ExperimentCheckpointManager.checkpoint` for where that happens.
"""
return Path(self._get_session_path(), "driver_artifacts").as_posix()
@property
def trial_fs_path(self) -> str:
"""The trial directory path on the `storage_filesystem`.
Raises a ValueError if `trial_dir_name` is not set beforehand.
"""
if self.trial_dir_name is None:
raise RuntimeError(
"Should not access `trial_fs_path` without setting `trial_dir_name`"
)
return Path(self.experiment_fs_path, self.trial_dir_name).as_posix()
@property
def trial_driver_staging_path(self) -> str:
"""The local filesystem path of the trial directory on the driver.
The driver is the node where `Trainer.fit`/`Tuner.fit` is being called.
This path is of the form:
`/tmp/ray/session_<session_id>/artifacts/<ray-train-job-timestamp>/
<experiment_dir_name>/driver_artifacts/<trial_dir_name>`
This should be used as the temporary location for files on the driver
before persisting them to `trial_fs_path`.
For example, callbacks (e.g., JsonLoggerCallback) should write trial-specific
logfiles within this directory.
"""
if self.trial_dir_name is None:
raise RuntimeError(
"Should not access `trial_driver_staging_path` "
"without setting `trial_dir_name`"
)
return Path(self.experiment_driver_staging_path, self.trial_dir_name).as_posix()
@property
def trial_working_directory(self) -> str:
"""The local filesystem path to trial working directory.
This path is of the form:
`/tmp/ray/session_<session_id>/artifacts/<ray-train-job-timestamp>/
<experiment_dir_name>/working_dirs/<trial_dir_name>`
Ray Train/Tune moves the remote actor's working directory to this path
by default, unless disabled by `RAY_CHDIR_TO_TRIAL_DIR` environment variable.
Writing files to this directory allows users to persist training artifacts
if `SyncConfig(sync_artifacts=True)` is set.
"""
if self.trial_dir_name is None:
raise RuntimeError(
"Cannot access `trial_working_directory` without "
"setting `trial_dir_name`"
)
return Path(
self._get_session_path(), "working_dirs", self.trial_dir_name
).as_posix()
@property
def checkpoint_fs_path(self) -> str:
"""The current checkpoint directory path on the `storage_filesystem`.
"Current" refers to the checkpoint that is currently being created/persisted.
The user of this class is responsible for setting the `current_checkpoint_index`
(e.g., incrementing when needed).
"""
return Path(self.trial_fs_path, self.checkpoint_dir_name).as_posix()
@property
def checkpoint_dir_name(self) -> str:
"""The current checkpoint directory name, based on the checkpoint index."""
return StorageContext._make_checkpoint_dir_name(self.current_checkpoint_index)
@staticmethod
def get_experiment_dir_name(run_obj: Union[str, Callable, Type]) -> str:
from ray.tune.experiment import Experiment
from ray.tune.utils import date_str
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_checkpoint_dir_name(index: int):
"""Get the name of the checkpoint directory, given an index."""
return f"checkpoint_{index:06d}"
python/ray/train/_internal/syncer.py
+429
View File
@@ -0,0 +1,429 @@
import abc
import logging
import threading
import time
import traceback
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple
from ray._private.thirdparty.tabulate.tabulate import tabulate
from ray.util.annotations import Deprecated, DeveloperAPI
from ray.widgets import Template
logger = logging.getLogger(__name__)
# Syncing period for syncing checkpoints between nodes or to cloud.
DEFAULT_SYNC_PERIOD = 300
# Default sync timeout after which syncing processes are aborted
DEFAULT_SYNC_TIMEOUT = 1800
@Deprecated
@dataclass
class SyncConfig:
sync_period: int = DEFAULT_SYNC_PERIOD
sync_timeout: int = DEFAULT_SYNC_TIMEOUT
sync_artifacts: bool = False
sync_artifacts_on_checkpoint: bool = True
def _repr_html_(self) -> str:
"""Generate an HTML representation of the SyncConfig."""
return Template("scrollableTable.html.j2").render(
table=tabulate(
{
"Setting": ["Sync period", "Sync timeout"],
"Value": [self.sync_period, self.sync_timeout],
},
tablefmt="html",
showindex=False,
headers="keys",
),
max_height="none",
)
class _BackgroundProcess:
def __init__(self, fn: Callable):
self._fn = fn
self._process = None
self._result = {}
self._start_time = float("-inf")
@property
def is_running(self):
return self._process and self._process.is_alive()
@property
def start_time(self):
return self._start_time
def start(self, *args, **kwargs):
if self.is_running:
return False
self._result = {}
def entrypoint():
try:
result = self._fn(*args, **kwargs)
except Exception as e:
self._result["exception"] = e
return
self._result["result"] = result
self._process = threading.Thread(target=entrypoint)
self._process.daemon = True
self._process.start()
self._start_time = time.time()
def wait(self, timeout: Optional[float] = None) -> Any:
"""Waits for the background process to finish running. Waits until the
background process has run for at least `timeout` seconds, counting from
the time when the process was started."""
if not self._process:
return None
time_remaining = None
if timeout:
elapsed = time.time() - self.start_time
time_remaining = max(timeout - elapsed, 0)
self._process.join(timeout=time_remaining)
if self._process.is_alive():
self._process = None
raise TimeoutError(
f"{getattr(self._fn, '__name__', str(self._fn))} did not finish "
f"running within the timeout of {timeout} seconds."
)
self._process = None
exception = self._result.get("exception")
if exception:
raise exception
result = self._result.get("result")
self._result = {}
return result
@DeveloperAPI
class Syncer(abc.ABC):
"""Syncer class for synchronizing data between Ray nodes and remote (cloud) storage.
This class handles data transfer for two cases:
1. Synchronizing data such as experiment state snapshots from the driver to
cloud storage.
2. Synchronizing data such as trial checkpoints from remote trainables to
cloud storage.
Synchronizing tasks are usually asynchronous and can be awaited using ``wait()``.
The base class implements a ``wait_or_retry()`` API that will retry a failed
sync command.
The base class also exposes an API to only kick off syncs every ``sync_period``
seconds.
Args:
sync_period: The minimum time in seconds between sync operations, as
used by ``sync_up/down_if_needed``.
sync_timeout: The maximum time to wait for a sync process to finish before
issuing a new sync operation. Ex: should be used by ``wait`` if launching
asynchronous sync tasks.
"""
def __init__(
self,
sync_period: float = DEFAULT_SYNC_PERIOD,
sync_timeout: float = DEFAULT_SYNC_TIMEOUT,
):
self.sync_period = sync_period
self.sync_timeout = sync_timeout
self.last_sync_up_time = float("-inf")
self.last_sync_down_time = float("-inf")
@abc.abstractmethod
def sync_up(
self, local_dir: str, remote_dir: str, exclude: Optional[List] = None
) -> bool:
"""Synchronize local directory to remote directory.
This function can spawn an asynchronous process that can be awaited in
``wait()``.
Args:
local_dir: Local directory to sync from.
remote_dir: Remote directory to sync up to. This is an URI
(``protocol://remote/path``).
exclude: Pattern of files to exclude, e.g.
``["*/checkpoint_*]`` to exclude trial checkpoints.
Returns:
True if sync process has been spawned, False otherwise.
"""
raise NotImplementedError
@abc.abstractmethod
def sync_down(
self, remote_dir: str, local_dir: str, exclude: Optional[List] = None
) -> bool:
"""Synchronize remote directory to local directory.
This function can spawn an asynchronous process that can be awaited in
``wait()``.
Args:
remote_dir: Remote directory to sync down from. This is an URI
(``protocol://remote/path``).
local_dir: Local directory to sync to.
exclude: Pattern of files to exclude, e.g.
``["*/checkpoint_*]`` to exclude trial checkpoints.
Returns:
True if sync process has been spawned, False otherwise.
"""
raise NotImplementedError
@abc.abstractmethod
def delete(self, remote_dir: str) -> bool:
"""Delete directory on remote storage.
This function can spawn an asynchronous process that can be awaited in
``wait()``.
Args:
remote_dir: Remote directory to delete. This is an URI
(``protocol://remote/path``).
Returns:
True if sync process has been spawned, False otherwise.
"""
raise NotImplementedError
def retry(self):
"""Retry the last sync up, sync down, or delete command.
You should implement this method if you spawn asynchronous syncing
processes.
"""
pass
def wait(self, timeout: Optional[float] = None):
"""Wait for asynchronous sync command to finish.
You should implement this method if you spawn asynchronous syncing
processes. This method should timeout after the asynchronous command
has run for `sync_timeout` seconds and raise a `TimeoutError`.
"""
pass
def sync_up_if_needed(
self, local_dir: str, remote_dir: str, exclude: Optional[List] = None
) -> bool:
"""Syncs up if time since last sync up is greater than sync_period.
Args:
local_dir: Local directory to sync from.
remote_dir: Remote directory to sync up to. This is an URI
(``protocol://remote/path``).
exclude: Pattern of files to exclude, e.g.
``["*/checkpoint_*]`` to exclude trial checkpoints.
Returns:
The result of ``sync_up`` if a sync was triggered, otherwise ``None``.
"""
now = time.time()
if now - self.last_sync_up_time >= self.sync_period:
result = self.sync_up(
local_dir=local_dir, remote_dir=remote_dir, exclude=exclude
)
self.last_sync_up_time = now
return result
def sync_down_if_needed(
self, remote_dir: str, local_dir: str, exclude: Optional[List] = None
):
"""Syncs down if time since last sync down is greater than sync_period.
Args:
remote_dir: Remote directory to sync down from. This is an URI
(``protocol://remote/path``).
local_dir: Local directory to sync to.
exclude: Pattern of files to exclude, e.g.
``["*/checkpoint_*]`` to exclude trial checkpoints.
Returns:
The result of ``sync_down`` if a sync was triggered, otherwise ``None``.
"""
now = time.time()
if now - self.last_sync_down_time >= self.sync_period:
result = self.sync_down(
remote_dir=remote_dir, local_dir=local_dir, exclude=exclude
)
self.last_sync_down_time = now
return result
def wait_or_retry(self, max_retries: int = 2, backoff_s: int = 5):
assert max_retries > 0
last_error_traceback = None
for i in range(max_retries + 1):
try:
self.wait()
except Exception as e:
attempts_remaining = max_retries - i
# If we're out of retries, then save the full traceback of the last
# error and show it when raising an exception.
if attempts_remaining == 0:
last_error_traceback = traceback.format_exc()
break
logger.error(
f"The latest sync operation failed with the following error: "
f"{repr(e)}\n"
f"Retrying {attempts_remaining} more time(s) after sleeping "
f"for {backoff_s} seconds..."
)
time.sleep(backoff_s)
self.retry()
continue
# Succeeded!
return
raise RuntimeError(
f"Failed sync even after {max_retries} retries. "
f"The latest sync failed with the following error:\n{last_error_traceback}"
)
def reset(self):
self.last_sync_up_time = float("-inf")
self.last_sync_down_time = float("-inf")
def close(self):
pass
def _repr_html_(self) -> str:
return
class _BackgroundSyncer(Syncer):
"""Syncer using a background process for asynchronous file transfer."""
def __init__(
self,
sync_period: float = DEFAULT_SYNC_PERIOD,
sync_timeout: float = DEFAULT_SYNC_TIMEOUT,
):
super(_BackgroundSyncer, self).__init__(
sync_period=sync_period, sync_timeout=sync_timeout
)
self._sync_process = None
self._current_cmd = None
def _should_continue_existing_sync(self):
"""Returns whether a previous sync is still running within the timeout."""
return (
self._sync_process
and self._sync_process.is_running
and time.time() - self._sync_process.start_time < self.sync_timeout
)
def _launch_sync_process(self, sync_command: Tuple[Callable, Dict]):
"""Waits for the previous sync process to finish,
then launches a new process that runs the given command."""
if self._sync_process:
try:
self.wait()
except Exception:
logger.warning(
f"Last sync command failed with the following error:\n"
f"{traceback.format_exc()}"
)
self._current_cmd = sync_command
self.retry()
def sync_up(
self, local_dir: str, remote_dir: str, exclude: Optional[List] = None
) -> bool:
if self._should_continue_existing_sync():
logger.debug(
f"Last sync still in progress, "
f"skipping sync up of {local_dir} to {remote_dir}"
)
return False
sync_up_cmd = self._sync_up_command(
local_path=local_dir, uri=remote_dir, exclude=exclude
)
self._launch_sync_process(sync_up_cmd)
return True
def _sync_up_command(
self, local_path: str, uri: str, exclude: Optional[List] = None
) -> Tuple[Callable, Dict]:
raise NotImplementedError
def sync_down(
self, remote_dir: str, local_dir: str, exclude: Optional[List] = None
) -> bool:
if self._should_continue_existing_sync():
logger.warning(
f"Last sync still in progress, "
f"skipping sync down of {remote_dir} to {local_dir}"
)
return False
sync_down_cmd = self._sync_down_command(uri=remote_dir, local_path=local_dir)
self._launch_sync_process(sync_down_cmd)
return True
def _sync_down_command(self, uri: str, local_path: str) -> Tuple[Callable, Dict]:
raise NotImplementedError
def delete(self, remote_dir: str) -> bool:
if self._should_continue_existing_sync():
logger.warning(
f"Last sync still in progress, skipping deletion of {remote_dir}"
)
return False
delete_cmd = self._delete_command(uri=remote_dir)
self._launch_sync_process(delete_cmd)
return True
def _delete_command(self, uri: str) -> Tuple[Callable, Dict]:
raise NotImplementedError
def wait(self, timeout: Optional[float] = None):
if self._sync_process:
try:
self._sync_process.wait(timeout=timeout or self.sync_timeout)
except Exception as e:
raise e
finally:
# Regardless of whether the sync process succeeded within the timeout,
# clear the sync process so a new one can be created.
self._sync_process = None
def retry(self):
if not self._current_cmd:
raise RuntimeError("No sync command set, cannot retry.")
cmd, kwargs = self._current_cmd
self._sync_process = _BackgroundProcess(cmd)
self._sync_process.start(**kwargs)
def __getstate__(self):
state = self.__dict__.copy()
state["_sync_process"] = None
return state
python/ray/train/_internal/utils.py
+230
View File
@@ -0,0 +1,230 @@
import abc
import functools
import inspect
import logging
import os
import socket
from typing import (
Any,
Callable,
ContextManager,
Dict,
List,
Optional,
Tuple,
TypeVar,
Union,
)
import ray
from ray._common.network_utils import find_free_port, is_ipv6
from ray.actor import ActorHandle
from ray.air._internal.util import (
StartTraceback,
StartTracebackWithWorkerRank,
)
from ray.exceptions import RayActorError
from ray.types import ObjectRef
T = TypeVar("T")
logger = logging.getLogger(__name__)
def check_for_failure(
remote_values: List[ObjectRef],
) -> Tuple[bool, Optional[Exception]]:
"""Check for actor failure when retrieving the remote values.
Args:
remote_values: List of object references from Ray actor methods.
Returns:
A tuple of (bool, Exception). The bool is
True if evaluating all object references is successful, False otherwise.
"""
unfinished = remote_values.copy()
while len(unfinished) > 0:
finished, unfinished = ray.wait(unfinished)
# If a failure occurs the ObjectRef will be marked as finished.
# Calling ray.get will expose the failure as a RayActorError.
for object_ref in finished:
# Everything in finished has either failed or completed
# successfully.
try:
ray.get(object_ref)
except RayActorError as exc:
failed_actor_rank = remote_values.index(object_ref)
logger.info(f"Worker {failed_actor_rank} has failed.")
return False, exc
except Exception as exc:
# Other (e.g. training) errors should be directly raised
failed_worker_rank = remote_values.index(object_ref)
raise StartTracebackWithWorkerRank(
worker_rank=failed_worker_rank
) from exc
return True, None
def get_address_and_port() -> Tuple[str, int]:
"""Returns the IP address and a free port on this node."""
addr = ray.util.get_node_ip_address()
port = find_free_port(socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET)
return addr, port
def update_env_vars(env_vars: Dict[str, Any]):
"""Updates the environment variables on this worker process.
Args:
env_vars: Environment variables to set.
"""
sanitized = {k: str(v) for k, v in env_vars.items()}
os.environ.update(sanitized)
def count_required_parameters(fn: Callable) -> int:
"""Counts the number of required parameters of a function.
NOTE: *args counts as 1 required parameter.
Args:
fn: The function whose required parameters should be counted.
Returns:
The number of required parameters of ``fn``.
Examples:
>>> def fn(a, b, /, c, *args, d=1, e=2, **kwargs):
... pass
>>> count_required_parameters(fn)
4
>>> fn = lambda: 1
>>> count_required_parameters(fn)
0
>>> def fn(config, a, b=1, c=2):
... pass
>>> from functools import partial
>>> count_required_parameters(partial(fn, a=0))
1
"""
params = inspect.signature(fn).parameters.values()
positional_param_kinds = {
inspect.Parameter.POSITIONAL_ONLY,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
inspect.Parameter.VAR_POSITIONAL,
}
return len(
[
p
for p in params
if p.default == inspect.Parameter.empty and p.kind in positional_param_kinds
]
)
def construct_train_func(
train_func: Union[Callable[[], T], Callable[[Dict[str, Any]], T]],
config: Optional[Dict[str, Any]],
train_func_context: ContextManager,
fn_arg_name: Optional[str] = "train_func",
discard_returns: bool = False,
) -> Callable[[], T]:
"""Validates and constructs the training function to execute.
Args:
train_func: The training function to execute.
This can either take in no arguments or a ``config`` dict.
config: Configurations to pass into ``train_func``. If None then an empty
Dict will be created.
train_func_context: Context manager for user's `train_func`, which executes
backend-specific logic before and after the training function.
fn_arg_name: The name of training function to use for error messages.
discard_returns: Whether to discard any returns from train_func or not.
Returns:
A valid training function.
Raises:
ValueError: if the input ``train_func`` is invalid.
"""
num_required_params = count_required_parameters(train_func)
if discard_returns:
# Discard any returns from the function so that
# BackendExecutor doesn't try to deserialize them.
# Those returns are inaccesible with AIR anyway.
@functools.wraps(train_func)
def discard_return_wrapper(*args, **kwargs):
try:
train_func(*args, **kwargs)
except Exception as e:
raise StartTraceback from e
wrapped_train_func = discard_return_wrapper
else:
wrapped_train_func = train_func
if num_required_params > 1:
err_msg = (
f"{fn_arg_name} should take in 0 or 1 required arguments, but it accepts "
f"{num_required_params} required arguments instead."
)
raise ValueError(err_msg)
elif num_required_params == 1:
config = {} if config is None else config
@functools.wraps(wrapped_train_func)
def train_fn():
try:
with train_func_context():
return wrapped_train_func(config)
except Exception as e:
raise StartTraceback from e
else: # num_params == 0
@functools.wraps(wrapped_train_func)
def train_fn():
try:
with train_func_context():
return wrapped_train_func()
except Exception as e:
raise StartTraceback from e
return train_fn
class Singleton(abc.ABCMeta):
"""Singleton Abstract Base Class
https://stackoverflow.com/questions/33364070/implementing
-singleton-as-metaclass-but-for-abstract-classes
"""
_instances = {}
def __call__(cls, *args, **kwargs):
if cls not in cls._instances:
cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
return cls._instances[cls]
class ActorWrapper:
"""Wraps an actor to provide same API as using the base class directly."""
def __init__(self, actor: ActorHandle):
self.actor = actor
def __getattr__(self, item):
# The below will fail if trying to access an attribute (not a method) from the
# actor.
actor_method = getattr(self.actor, item)
return lambda *args, **kwargs: ray.get(actor_method.remote(*args, **kwargs))
python/ray/train/_internal/worker_group.py
+445
View File
@@ -0,0 +1,445 @@
import copy
import logging
import os
import socket
from collections import defaultdict
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
import ray
from ray.actor import ActorHandle
from ray.air._internal.util import exception_cause, skip_exceptions
from ray.train._internal.base_worker_group import BaseWorkerGroup
from ray.types import ObjectRef
from ray.util.placement_group import PlacementGroup
T = TypeVar("T")
logger = logging.getLogger(__name__)
class RayTrainWorker:
"""A class to execute arbitrary functions. Does not hold any state."""
def __execute(self, func: Callable[..., T], *args, **kwargs) -> T:
"""Executes the input function and returns the output.
Args:
func: The function to execute.
*args: Positional arguments to pass into ``func``.
**kwargs: Keyword arguments to pass into ``func``.
Returns:
The result of calling ``func`` with the provided arguments.
"""
try:
return func(*args, **kwargs)
except Exception as e:
skipped = skip_exceptions(e)
raise skipped from exception_cause(skipped)
@dataclass
class WorkerMetadata:
"""Metadata for each worker/actor.
This information is expected to stay the same throughout the lifetime of
actor.
Args:
node_id: ID of the node this worker is on.
node_ip: IP address of the node this worker is on.
hostname: Hostname that this worker is on.
resource_ids: Map of accelerator resources
("GPU", "neuron_cores", ..) to their IDs.
pid: Process ID of this worker.
"""
node_id: str
node_ip: str
hostname: str
resource_ids: Dict[str, List[str]]
pid: int
@dataclass
class Worker:
"""Class representing a Worker."""
actor: ActorHandle
metadata: WorkerMetadata
def create_executable_class(executable_cls: Optional[Type] = None) -> Type:
"""Create the executable class to use as the Ray actors."""
if not executable_cls:
return RayTrainWorker
elif issubclass(executable_cls, RayTrainWorker):
return executable_cls
else:
class _WrappedExecutable(executable_cls, RayTrainWorker):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
return _WrappedExecutable
def construct_metadata() -> WorkerMetadata:
"""Creates metadata for this worker.
This function is expected to be run on the actor.
"""
node_id = ray.get_runtime_context().get_node_id()
node_ip = ray.util.get_node_ip_address()
hostname = socket.gethostname()
accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
pid = os.getpid()
return WorkerMetadata(
node_id=node_id,
node_ip=node_ip,
hostname=hostname,
resource_ids=accelerator_ids,
pid=pid,
)
class WorkerGroup(BaseWorkerGroup):
"""Group of Ray Actors that can execute arbitrary functions.
``WorkerGroup`` launches Ray actors according to the given
specification. It can then execute arbitrary Python functions in each of
these workers.
If not enough resources are available to launch the actors, the Ray
cluster will automatically scale up if autoscaling is enabled.
Args:
num_workers: The number of workers (Ray actors) to launch.
Defaults to 1.
resources_per_worker: Dictionary specifying the resources that will be
requested for each worker. Defaults to {"CPU": 1}.
actor_cls: If specified use this class as the remote actors.
actor_cls_args: If ``actor_cls`` is provided, these positional args will
be used for the worker initialization.
actor_cls_kwargs: If ``actor_cls`` is provided, these keyword args will
be used for the worker initialization.
placement_group: The placement group that workers
should be created in. Defaults to "default" which will inherit the
parent placement group (if child tasks should be captured).
Example:
.. code_block:: python
worker_group = WorkerGroup(num_workers=2)
output = worker_group.execute(lambda: 1)
assert len(output) == 2
assert all(o == 1 for o in output)
"""
def __init__(
self,
num_workers: int = 1,
resources_per_worker: Optional[Dict[str, float]] = None,
actor_cls: Type = None,
actor_cls_args: Optional[Tuple] = None,
actor_cls_kwargs: Optional[Dict] = None,
placement_group: Union[PlacementGroup, str] = "default",
):
if resources_per_worker is None:
resources_per_worker = {"CPU": 1}
else:
resources_per_worker = resources_per_worker.copy()
if num_workers <= 0:
raise ValueError(
"The provided `num_workers` must be greater "
f"than 0. Received num_workers={num_workers} "
f"instead."
)
if any(v < 0 for v in resources_per_worker.values()):
raise ValueError(
"The number of resources per worker must not be negative. "
f"Received resources_per_worker={resources_per_worker}."
)
if (actor_cls_args or actor_cls_kwargs) and not actor_cls:
raise ValueError(
"`actor_cls_args` or `actor_class_kwargs` are "
"passed in but no `actor_cls` is passed in."
)
self.num_workers = num_workers
self.resources_per_worker = resources_per_worker
_resources_per_worker = copy.deepcopy(resources_per_worker)
self.num_cpus_per_worker = _resources_per_worker.pop("CPU", 0)
self.num_gpus_per_worker = _resources_per_worker.pop("GPU", 0)
self.memory_per_worker = _resources_per_worker.pop("memory", 0)
self.workers = []
self._base_cls = create_executable_class(actor_cls)
assert issubclass(self._base_cls, RayTrainWorker)
self._actor_cls_args = actor_cls_args or []
self._actor_cls_kwargs = actor_cls_kwargs or {}
self._placement_group = placement_group
# TODO(matt): Validate resources. Fast-fail if it is impossible to
# handle the request, rather than hang indefinitely.
self._remote_cls = ray.remote(
num_cpus=self.num_cpus_per_worker,
num_gpus=self.num_gpus_per_worker,
memory=self.memory_per_worker,
resources=_resources_per_worker,
)(self._base_cls)
self.start()
def start(self):
"""Starts all the workers in this worker group."""
if self.workers and len(self.workers) > 0:
raise RuntimeError(
"The workers have already been started. "
"Please call `shutdown` first if you want to "
"restart them."
)
logger.debug(f"Starting {self.num_workers} workers.")
self.add_workers(self.num_workers)
logger.debug(f"{len(self.workers)} workers have successfully started.")
def shutdown(self, patience_s: float = 5):
"""Shutdown all the workers in this worker group.
Args:
patience_s: Attempt a graceful shutdown
of the workers for this many seconds. Fallback to force kill
if graceful shutdown is not complete after this time. If
this is less than or equal to 0, immediately force kill all
workers.
"""
logger.debug(f"Shutting down {len(self.workers)} workers.")
if patience_s <= 0:
for worker in self.workers:
ray.kill(worker.actor)
else:
done_refs = [w.actor.__ray_terminate__.remote() for w in self.workers]
# Wait for actors to die gracefully.
done, not_done = ray.wait(done_refs, timeout=patience_s)
if not_done:
logger.debug("Graceful termination failed. Falling back to force kill.")
# If all actors are not able to die gracefully, then kill them.
for worker in self.workers:
ray.kill(worker.actor)
logger.debug("Shutdown successful.")
self.workers = []
def execute_async(self, func: Callable[..., T], *args, **kwargs) -> List[ObjectRef]:
"""Execute ``func`` on each worker and return the futures.
Args:
func: A function to call on each worker.
*args: Positional arguments passed directly into ``func``.
**kwargs: Keyword arguments passed directly into ``func``.
Returns:
(List[ObjectRef]) A list of ``ObjectRef`` representing the
output of ``func`` from each worker. The order is the same
as ``self.workers``.
"""
if len(self.workers) <= 0:
raise RuntimeError(
"There are no active workers. This worker "
"group has most likely been shut down. Please"
"create a new WorkerGroup or restart this one."
)
return [
w.actor._RayTrainWorker__execute.options(
name=f"_RayTrainWorker__execute.{func.__name__}"
).remote(func, *args, **kwargs)
for w in self.workers
]
def execute(self, func: Callable[..., T], *args, **kwargs) -> List[T]:
"""Execute ``func`` on each worker and return the outputs of ``func``.
Args:
func: A function to call on each worker.
*args: Positional arguments passed directly into ``func``.
**kwargs: Keyword arguments passed directly into ``func``.
Returns:
(List[T]) A list containing the output of ``func`` from each
worker. The order is the same as ``self.workers``.
"""
# TODO: Add a timeout in the case of a hang, particularly
# relevant when func is TorchConfig.on_shutdown
return ray.get(self.execute_async(func, *args, **kwargs))
def execute_single_async(
self, worker_index: int, func: Callable[..., T], *args, **kwargs
) -> ObjectRef:
"""Execute ``func`` on worker ``worker_index`` and return futures.
Args:
worker_index: The index to execute func on.
func: A function to call on the first worker.
*args: Positional arguments passed directly into ``func``.
**kwargs: Keyword arguments passed directly into ``func``.
Returns:
(ObjectRef) An ObjectRef representing the output of func.
"""
if worker_index >= len(self.workers):
raise ValueError(
f"The provided worker_index {worker_index} is "
f"not valid for {self.num_workers} workers."
)
return (
self.workers[worker_index]
.actor._RayTrainWorker__execute.options(
name=f"_RayTrainWorker__execute.{func.__name__}"
)
.remote(func, *args, **kwargs)
)
def execute_single(
self, worker_index: int, func: Callable[..., T], *args, **kwargs
) -> T:
"""Execute ``func`` on worker with index ``worker_index``.
Args:
worker_index: The index to execute func on.
func: A function to call on the first worker.
*args: Positional arguments passed directly into ``func``.
**kwargs: Keyword arguments passed directly into ``func``.
Returns:
(T) The output of func.
"""
return ray.get(self.execute_single_async(worker_index, func, *args, **kwargs))
def remove_workers(self, worker_indexes: List[int]):
"""Removes the workers with the specified indexes.
The removed workers will go out of scope and their actor processes
will be terminated.
Args:
worker_indexes: The indexes of the workers to remove.
"""
new_workers = []
for i in range(len(self.workers)):
if i not in worker_indexes:
new_workers.append(self.workers[i])
self.workers = new_workers
def add_workers(self, num_workers: int):
"""Adds ``num_workers`` to this WorkerGroup.
Note: Adding workers when the cluster/placement group is at capacity
may lead to undefined hanging behavior. If you are attempting to
replace existing workers in the WorkerGroup, remove_workers() should
be called first.
Args:
num_workers: The number of workers to add.
"""
new_actors = []
new_actor_metadata = []
for _ in range(num_workers):
actor = self._remote_cls.options(
placement_group=self._placement_group
).remote(*self._actor_cls_args, **self._actor_cls_kwargs)
new_actors.append(actor)
new_actor_metadata.append(
actor._RayTrainWorker__execute.options(
name="_RayTrainWorker__execute.construct_metadata"
).remote(construct_metadata)
)
# Get metadata from all actors.
metadata = ray.get(new_actor_metadata)
for i in range(len(new_actors)):
self.workers.append(Worker(actor=new_actors[i], metadata=metadata[i]))
def sort_workers_by_node_id_and_gpu_id(self, _first_node_id: Optional[str] = None):
"""Reorder the workers by their node id and the lowest GPU id.
This is useful for collocating workers on the same node.
Example:
Given workers with the following attributes:
worker_0: node_id=1, gpu_ids=[1]
worker_1: node_id=0, gpu_ids=[0]
worker_2: node_id=1, gpu_ids=[0]
worker_3: node_id=0, gpu_ids=[1]
The function will perform the following steps:
1. Group by node ID:
node_id=0: worker_1, worker_3
node_id=1: worker_0, worker_2
2. Sort each group by GPU ID:
node_id=0: worker_1 (gpu_id=0), worker_3 (gpu_id=1)
node_id=1: worker_2 (gpu_id=0), worker_0 (gpu_id=1)
Resulting in the order: [worker_1, worker_3, worker_2, worker_0]
Args:
_first_node_id: The first ID to group by.
Set this to the node ID of the trainer coordinator to ensure that the
rank 0 worker is on the same node, allowing additional resources to
be specified for rank 0 workers via
`ScalingConfig(trainer_resources=)`.
"""
node_id_to_workers = defaultdict(list)
if _first_node_id is not None:
node_id_to_workers[_first_node_id] = []
for worker in self.workers:
node_id_to_workers[worker.metadata.node_id].append(worker)
# Sort workers on the same node by the lowest GPU id
# More details: https://github.com/ray-project/ray/issues/40803
def get_lowest_gpu_id(worker) -> int:
gpu_ids = worker.metadata.resource_ids.get("GPU", [])
# If there are no GPU IDs, return 0 as a default
if not gpu_ids:
return 0
# Attempt to convert GPU IDs to integers and find the minimum ID.
# Fallback to return the minimum string-based ID
try:
return min(int(gpu_id) for gpu_id in gpu_ids)
except ValueError:
return min(gpu_ids)
for node_id in node_id_to_workers:
node_id_to_workers[node_id].sort(key=get_lowest_gpu_id)
sorted_workers = []
for workers in node_id_to_workers.values():
sorted_workers.extend(workers)
self.workers = sorted_workers
def __len__(self):
return len(self.workers)
def get_resources_per_worker(self) -> dict:
"""Get the resources allocated per worker."""
return copy.deepcopy(self.resources_per_worker)