chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,16 @@
|
||||
from .accelerators import AcceleratorSetupCallback
|
||||
from .backend_setup import BackendSetupCallback
|
||||
from .datasets import DatasetsCallback
|
||||
from .state_manager import StateManagerCallback
|
||||
from .working_dir_setup import WorkingDirectorySetupCallback
|
||||
|
||||
__all__ = [
|
||||
"AcceleratorSetupCallback",
|
||||
"BackendSetupCallback",
|
||||
"DatasetsCallback",
|
||||
"StateManagerCallback",
|
||||
"WorkingDirectorySetupCallback",
|
||||
]
|
||||
|
||||
|
||||
# DO NOT ADD ANYTHING AFTER THIS LINE.
|
||||
@@ -0,0 +1,160 @@
|
||||
import logging
|
||||
import os
|
||||
from collections import defaultdict
|
||||
from typing import TYPE_CHECKING, Any, Dict, List
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray._private.accelerators.nvidia_gpu import CUDA_VISIBLE_DEVICES_ENV_VAR
|
||||
from ray._private.ray_constants import env_bool
|
||||
from ray.train import BackendConfig
|
||||
from ray.train.constants import ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV
|
||||
from ray.train.v2._internal.execution.callback import WorkerGroupCallback
|
||||
from ray.train.v2._internal.execution.worker_group import ActorMetadata
|
||||
from ray.train.v2.api.config import ScalingConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.worker_group.worker import Worker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AcceleratorSetupCallback(WorkerGroupCallback):
|
||||
"""Perform accelerator setup for workers.
|
||||
|
||||
For example, this callback can be used to share CUDA_VISIBLE_DEVICES
|
||||
among workers on the same node.
|
||||
"""
|
||||
|
||||
def __init__(self, backend_config: BackendConfig, scaling_config: ScalingConfig):
|
||||
self._backend = backend_config.backend_cls()
|
||||
self._scaling_config = scaling_config
|
||||
|
||||
def before_init_train_context(
|
||||
self, workers: List["Worker"]
|
||||
) -> Dict[str, List[Any]]:
|
||||
self._maybe_share_cuda_visible_devices(workers)
|
||||
# TODO: Add support for sharing other accelerator resources.
|
||||
|
||||
return {}
|
||||
|
||||
def _maybe_share_cuda_visible_devices(self, workers: List["Worker"]):
|
||||
"""Set CUDA visible devices environment variables on workers."""
|
||||
share_cuda_visible_devices_enabled = env_bool(
|
||||
ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
|
||||
self._backend.share_cuda_visible_devices,
|
||||
)
|
||||
|
||||
if (
|
||||
self._scaling_config._resources_per_worker_not_none.get("GPU", 0) > 0
|
||||
and share_cuda_visible_devices_enabled
|
||||
):
|
||||
_share_cuda_visible_devices(workers)
|
||||
|
||||
|
||||
def _share_cuda_visible_devices(workers: List["Worker"]):
|
||||
"""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"
|
||||
|
||||
Args:
|
||||
workers: List of worker objects.
|
||||
"""
|
||||
_share_accelerator_ids(workers, ray_constants.GPU, CUDA_VISIBLE_DEVICES_ENV_VAR)
|
||||
|
||||
|
||||
def _share_accelerator_ids(
|
||||
workers: List["Worker"], accelerator_name: 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"
|
||||
- Worker3: "0,1"
|
||||
|
||||
Args:
|
||||
workers: List of worker objects.
|
||||
accelerator_name: The name of the accelerator.
|
||||
env_var: The name of the environment variable to set.
|
||||
"""
|
||||
worker_metadatas = [worker.metadata for worker in workers]
|
||||
visible_accelerator_ids_per_worker = _get_visible_accelerator_ids_per_worker(
|
||||
worker_metadatas=worker_metadatas, accelerator_name=accelerator_name
|
||||
)
|
||||
|
||||
def set_accelerator_ids(accelerator_ids):
|
||||
os.environ[env_var] = accelerator_ids
|
||||
|
||||
futures = []
|
||||
for rank, visible_accelerator_ids in enumerate(visible_accelerator_ids_per_worker):
|
||||
futures.append(
|
||||
workers[rank].execute_async(
|
||||
set_accelerator_ids, accelerator_ids=visible_accelerator_ids
|
||||
)
|
||||
)
|
||||
ray.get(futures)
|
||||
|
||||
|
||||
def _get_visible_accelerator_ids_per_worker(
|
||||
worker_metadatas: List[ActorMetadata], accelerator_name: str
|
||||
) -> List[str]:
|
||||
"""Returns a list of comma-separated accelerator IDs visible to each worker.
|
||||
|
||||
All workers on a node should have the same set of visible accelerators,
|
||||
which is the union of accelerator ids of the workers.
|
||||
|
||||
Args:
|
||||
worker_metadatas: The actor metadata for each worker.
|
||||
accelerator_name: The name of the accelerator resource to inspect.
|
||||
|
||||
Returns:
|
||||
A list of comma-separated accelerator ID strings. This list is the
|
||||
same length as the number of workers.
|
||||
"""
|
||||
for metadata in worker_metadatas:
|
||||
if accelerator_name not in metadata.accelerator_ids:
|
||||
raise ValueError(
|
||||
f"Accelerator '{accelerator_name}' is not available on all workers. "
|
||||
f"Got these available accelerators instead: {metadata.accelerator_ids}"
|
||||
)
|
||||
|
||||
node_id_to_accelerator_ids = defaultdict(set)
|
||||
|
||||
for metadata in worker_metadatas:
|
||||
node_id_to_accelerator_ids[metadata.node_id].update(
|
||||
metadata.accelerator_ids[accelerator_name]
|
||||
)
|
||||
|
||||
visible_accelerator_ids_per_worker = []
|
||||
for worker_id in range(len(worker_metadatas)):
|
||||
node_id = worker_metadatas[worker_id].node_id
|
||||
accelerator_ids = sorted(node_id_to_accelerator_ids[node_id])
|
||||
all_resource_ids = ",".join([str(id) for id in accelerator_ids])
|
||||
visible_accelerator_ids_per_worker.append(all_resource_ids)
|
||||
|
||||
return visible_accelerator_ids_per_worker
|
||||
@@ -0,0 +1,33 @@
|
||||
import logging
|
||||
|
||||
from ray.exceptions import RayActorError, RayTaskError
|
||||
from ray.train.backend import BackendConfig
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ReplicaGroupCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
ExecutionGroup,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BackendSetupCallback(ReplicaGroupCallback, WorkerGroupCallback):
|
||||
def __init__(self, backend_config: BackendConfig):
|
||||
self._backend_config = backend_config
|
||||
self._backend = backend_config.backend_cls()
|
||||
|
||||
def after_execution_group_start(self, execution_group: ExecutionGroup):
|
||||
self._backend.on_start(execution_group, self._backend_config)
|
||||
self._backend.on_training_start(execution_group, self._backend_config)
|
||||
|
||||
def before_execution_group_shutdown(self, execution_group: ExecutionGroup):
|
||||
try:
|
||||
self._backend.on_shutdown(execution_group, self._backend_config)
|
||||
except (RayActorError, RayTaskError):
|
||||
logger.warning(
|
||||
"Graceful shutdown of backend failed. This is "
|
||||
"expected if one of the workers has crashed.",
|
||||
exc_info=True,
|
||||
)
|
||||
@@ -0,0 +1,175 @@
|
||||
import copy
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional
|
||||
|
||||
import ray
|
||||
import ray.train
|
||||
from ray.train.v2._internal.data_integration.interfaces import (
|
||||
DatasetShardMetadata,
|
||||
DatasetShardProvider,
|
||||
GenDataset,
|
||||
)
|
||||
from ray.train.v2._internal.execution.callback import WorkerGroupCallback
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.worker_group.worker_group import (
|
||||
Worker,
|
||||
WorkerGroup,
|
||||
WorkerGroupContext,
|
||||
)
|
||||
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import DataIterator, Dataset, NodeIdStr
|
||||
from ray.data.context import DataContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RayDatasetShardProvider:
|
||||
def __init__(
|
||||
self,
|
||||
datasets: Dict[str, GenDataset],
|
||||
data_config: ray.train.DataConfig,
|
||||
data_context: "DataContext",
|
||||
world_size: int,
|
||||
worker_node_ids: List["NodeIdStr"],
|
||||
):
|
||||
from ray.train.v2._internal.data_integration.dataset_manager import (
|
||||
DatasetManager,
|
||||
)
|
||||
|
||||
self._dataset_names = set(datasets)
|
||||
self._dataset_manager = (
|
||||
ray.remote(DatasetManager)
|
||||
.options(
|
||||
num_cpus=0,
|
||||
scheduling_strategy=NodeAffinitySchedulingStrategy(
|
||||
ray.get_runtime_context().get_node_id(), soft=False
|
||||
),
|
||||
)
|
||||
.remote(
|
||||
datasets=datasets,
|
||||
data_config=data_config,
|
||||
data_context=data_context,
|
||||
world_size=world_size,
|
||||
worker_node_ids=worker_node_ids,
|
||||
)
|
||||
)
|
||||
self._cached_dataset_shards: Dict[str, "DataIterator"] = {}
|
||||
|
||||
def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
|
||||
dataset_name = dataset_info.dataset_name
|
||||
if dataset_name not in self._dataset_names:
|
||||
raise KeyError(
|
||||
f"Dataset shard for '{dataset_name}' not found. "
|
||||
"Please ensure that the dataset is passed through the Trainer `datasets` "
|
||||
"argument."
|
||||
)
|
||||
|
||||
if dataset_name not in self._cached_dataset_shards:
|
||||
self._cached_dataset_shards[dataset_name] = ray.get(
|
||||
self._dataset_manager.get_dataset_shard.remote(dataset_info)
|
||||
)
|
||||
|
||||
return self._cached_dataset_shards[dataset_name]
|
||||
|
||||
def shutdown_data_executors(self) -> None:
|
||||
"""
|
||||
Attempts to eagerly shutdown the data executors for datasets, freeing resources allocated to data execution.
|
||||
"""
|
||||
try:
|
||||
self._dataset_manager.shutdown_data_executors.remote()
|
||||
except Exception:
|
||||
logger.debug("Failed to invoke remote cleanup of Dataset Manager.")
|
||||
|
||||
|
||||
class DatasetsCallback(WorkerGroupCallback):
|
||||
"""A callback for managing Ray Datasets for the worker group."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_run_context: TrainRunContext,
|
||||
datasets: Dict[str, "Dataset"],
|
||||
):
|
||||
self._datasets = datasets
|
||||
self._data_config = copy.deepcopy(train_run_context.dataset_config)
|
||||
self._scaling_config = train_run_context.scaling_config
|
||||
self._dataset_shard_provider: Optional[RayDatasetShardProvider] = None
|
||||
|
||||
# Capture the current DataContext to propagate it to
|
||||
# the Train workers later.
|
||||
# The propagation works in the following way:
|
||||
# 1. This callback is created when user create the Trainer.
|
||||
# 2. Then this callback will be passed to the Controller actor.
|
||||
# 3. Lastly, when the worker group is initialized, the Controller
|
||||
# will call the `after_worker_group_start` callback to propagate
|
||||
# the DataContext to Train workers.
|
||||
from ray.data.context import DataContext
|
||||
|
||||
self._data_context = copy.deepcopy(DataContext.get_current())
|
||||
|
||||
def get_train_total_resources(
|
||||
self, scaling_config: ray.train.ScalingConfig
|
||||
) -> Dict[str, float]:
|
||||
"""Return the resources reserved for training, so that Data can exclude
|
||||
these resources logically from its available pool."""
|
||||
if scaling_config.elasticity_enabled:
|
||||
# If Train is running with a variable number of workers,
|
||||
# we can't provide a fixed number of resources to exclude.
|
||||
# Instead, Train and Data should coordinate via the autoscaling
|
||||
# coordinator to allocate resources dynamically.
|
||||
return {}
|
||||
return scaling_config.total_resources
|
||||
|
||||
# --------------------------
|
||||
# WorkerGroupCallback
|
||||
# --------------------------
|
||||
|
||||
def before_init_train_context(
|
||||
self, workers: List[Worker]
|
||||
) -> Dict[str, List[DatasetShardProvider]]:
|
||||
world_size = len(workers)
|
||||
worker_node_ids = [worker.metadata.node_id for worker in workers]
|
||||
datasets = {k: v() if callable(v) else v for k, v in self._datasets.items()}
|
||||
|
||||
# TODO: Move this to the constructor.
|
||||
# Notify the DataConfig about the total resources reserved for training.
|
||||
total_train_resources = self.get_train_total_resources(self._scaling_config)
|
||||
self._data_config.set_train_total_resources(
|
||||
total_train_resources.get("CPU", 0), total_train_resources.get("GPU", 0)
|
||||
)
|
||||
|
||||
self._dataset_shard_provider = RayDatasetShardProvider(
|
||||
datasets=datasets,
|
||||
data_config=self._data_config,
|
||||
data_context=self._data_context,
|
||||
world_size=world_size,
|
||||
worker_node_ids=worker_node_ids,
|
||||
)
|
||||
return {"dataset_shard_provider": [self._dataset_shard_provider] * world_size}
|
||||
|
||||
def after_worker_group_start(self, worker_group: WorkerGroup):
|
||||
# Propagate DataContext
|
||||
from ray.data.context import DataContext
|
||||
|
||||
def _propagate_data_context(ctx: "DataContext"):
|
||||
DataContext._set_current(ctx)
|
||||
|
||||
worker_group.execute(
|
||||
_propagate_data_context,
|
||||
self._data_context,
|
||||
)
|
||||
|
||||
def after_worker_group_shutdown(
|
||||
self, worker_group_context: WorkerGroupContext
|
||||
) -> None:
|
||||
shard_provider = self._dataset_shard_provider
|
||||
if shard_provider:
|
||||
shard_provider.shutdown_data_executors()
|
||||
|
||||
def after_worker_group_abort(
|
||||
self, worker_group_context: WorkerGroupContext
|
||||
) -> None:
|
||||
shard_provider = self._dataset_shard_provider
|
||||
if shard_provider:
|
||||
shard_provider.shutdown_data_executors()
|
||||
@@ -0,0 +1,39 @@
|
||||
import importlib
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
from ray.train.v2._internal.constants import RAY_TRAIN_CALLBACKS_ENV_VAR
|
||||
from ray.train.v2._internal.execution.callback import RayTrainCallback
|
||||
|
||||
|
||||
def _initialize_env_callbacks() -> List[RayTrainCallback]:
|
||||
"""Initialize callbacks from environment variable.
|
||||
|
||||
Returns:
|
||||
List of callbacks initialized from environment variable.
|
||||
"""
|
||||
callbacks = []
|
||||
callbacks_str = os.environ.get(RAY_TRAIN_CALLBACKS_ENV_VAR, "")
|
||||
if not callbacks_str:
|
||||
return callbacks
|
||||
|
||||
for callback_path in callbacks_str.split(","):
|
||||
callback_path = callback_path.strip()
|
||||
if not callback_path:
|
||||
continue
|
||||
|
||||
try:
|
||||
module_path, class_name = callback_path.rsplit(".", 1)
|
||||
module = importlib.import_module(module_path)
|
||||
callback_cls = getattr(module, class_name)
|
||||
if not issubclass(callback_cls, RayTrainCallback):
|
||||
raise TypeError(
|
||||
f"Callback class '{callback_path}' must be a subclass of "
|
||||
f"RayTrainCallback, got {type(callback_cls).__name__}"
|
||||
)
|
||||
callback = callback_cls()
|
||||
callbacks.append(callback)
|
||||
except (ImportError, AttributeError, ValueError, TypeError) as e:
|
||||
raise ValueError(f"Failed to import callback from '{callback_path}'") from e
|
||||
|
||||
return callbacks
|
||||
@@ -0,0 +1,124 @@
|
||||
from contextlib import contextmanager
|
||||
from typing import Dict, Optional
|
||||
|
||||
import ray
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ControllerCallback,
|
||||
TrainContextCallback,
|
||||
WorkerCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext, get_train_context
|
||||
from ray.train.v2._internal.execution.controller.state import (
|
||||
TrainControllerState,
|
||||
TrainControllerStateType,
|
||||
)
|
||||
from ray.train.v2._internal.metrics.base import Metric
|
||||
from ray.train.v2._internal.metrics.controller import ControllerMetrics
|
||||
from ray.train.v2._internal.metrics.worker import WorkerMetrics
|
||||
from ray.train.v2._internal.util import time_monotonic
|
||||
|
||||
|
||||
class ControllerMetricsCallback(ControllerCallback, WorkerGroupCallback):
|
||||
"""Callback that records controller-specific metrics."""
|
||||
|
||||
def after_controller_start(self, train_run_context: TrainRunContext):
|
||||
"""Initialize metrics after controller starts."""
|
||||
self._run_name = train_run_context.get_run_config().name
|
||||
self._run_id = train_run_context.run_id
|
||||
self._metrics: Dict[str, Metric] = ControllerMetrics.get_controller_metrics(
|
||||
self._run_name, self._run_id
|
||||
)
|
||||
# Record initial state
|
||||
self._metrics[ControllerMetrics.CONTROLLER_STATE].record(
|
||||
TrainControllerStateType.INITIALIZING
|
||||
)
|
||||
|
||||
async def before_controller_shutdown(self):
|
||||
"""Shutdown metrics before controller shuts down."""
|
||||
for metric in self._metrics.values():
|
||||
metric.reset()
|
||||
|
||||
def after_controller_state_update(
|
||||
self,
|
||||
previous_state: TrainControllerState,
|
||||
current_state: TrainControllerState,
|
||||
):
|
||||
"""Record state transitions after controller state updates."""
|
||||
self._metrics[ControllerMetrics.CONTROLLER_STATE].record(
|
||||
current_state._state_type
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def on_worker_group_start(self):
|
||||
"""Measure time taken to start worker group."""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[ControllerMetrics.WORKER_GROUP_START_TOTAL_TIME_S].record(
|
||||
elapsed_time_s
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def on_worker_group_shutdown(self):
|
||||
"""Measure time taken to shutdown worker group."""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[ControllerMetrics.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S].record(
|
||||
elapsed_time_s
|
||||
)
|
||||
|
||||
|
||||
class WorkerMetricsCallback(WorkerCallback, TrainContextCallback):
|
||||
"""Callback that records worker-specific metrics."""
|
||||
|
||||
def __init__(self, train_run_context: TrainRunContext):
|
||||
self._run_name = train_run_context.get_run_config().name
|
||||
self._run_id = train_run_context.run_id
|
||||
self._metrics: Optional[Dict[str, Metric]] = None
|
||||
|
||||
def after_init_train_context(self):
|
||||
"""Initialize metrics after train context is initialized."""
|
||||
train_context = get_train_context()
|
||||
core_context = ray.runtime_context.get_runtime_context()
|
||||
world_rank = train_context.get_world_rank()
|
||||
worker_actor_id = core_context.get_actor_id()
|
||||
self._metrics = WorkerMetrics.get_worker_metrics(
|
||||
self._run_name, self._run_id, world_rank, worker_actor_id
|
||||
)
|
||||
|
||||
def before_worker_shutdown(self):
|
||||
"""Shutdown metrics before shutdown."""
|
||||
if self._metrics:
|
||||
for metric in self._metrics.values():
|
||||
metric.reset()
|
||||
|
||||
@contextmanager
|
||||
def on_report(self):
|
||||
"""
|
||||
Context manager to measure the time taken to report a checkpoint to the storage.
|
||||
"""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[WorkerMetrics.REPORT_TOTAL_BLOCKED_TIME_S].record(elapsed_time_s)
|
||||
|
||||
@contextmanager
|
||||
def on_checkpoint_sync(self):
|
||||
"""Measure time spent in the cross-rank barrier that synchronizes the
|
||||
checkpoint directory name across all workers."""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[WorkerMetrics.CHECKPOINT_SYNC_TOTAL_TIME_S].record(elapsed_time_s)
|
||||
|
||||
@contextmanager
|
||||
def on_checkpoint_transfer(self):
|
||||
"""Measure time spent transferring checkpoint files to storage."""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[WorkerMetrics.CHECKPOINT_TRANSFER_TOTAL_TIME_S].record(
|
||||
elapsed_time_s
|
||||
)
|
||||
@@ -0,0 +1,131 @@
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import ray
|
||||
from ray.exceptions import RayActorError
|
||||
from ray.train.v2._internal.constants import GET_ACTOR_TIMEOUT_S
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ControllerCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.controller.placement_group_cleaner import (
|
||||
PlacementGroupCleaner,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.worker_group import WorkerGroup
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PlacementGroupCleanerCallback(ControllerCallback, WorkerGroupCallback):
|
||||
"""Callback that manages a PlacementGroupCleaner for the training controller.
|
||||
|
||||
This callback ensures that placement groups are cleaned up even if the controller
|
||||
dies ungracefully.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
check_interval_s: float = 1.0,
|
||||
get_actor_timeout_s: float = GET_ACTOR_TIMEOUT_S,
|
||||
stop_timeout: Optional[float] = None,
|
||||
):
|
||||
"""Initialize the callback.
|
||||
|
||||
Args:
|
||||
check_interval_s: How often (in seconds) the cleaner should check
|
||||
if the controller is still alive.
|
||||
get_actor_timeout_s: How long to wait when calling the get actor state api.
|
||||
stop_timeout: How long to wait for the cleaner to stop.
|
||||
"""
|
||||
self._check_interval_s = check_interval_s
|
||||
self._get_actor_timeout_s = get_actor_timeout_s
|
||||
self._stop_timeout = stop_timeout
|
||||
if self._stop_timeout is None:
|
||||
self._stop_timeout = max(
|
||||
2.0, self._check_interval_s * 2 + self._get_actor_timeout_s
|
||||
)
|
||||
self._cleaner: Optional[PlacementGroupCleaner] = None
|
||||
self._controller_actor_id: Optional[str] = None
|
||||
|
||||
def after_controller_start(self, train_run_context: "TrainRunContext"):
|
||||
"""Launch the detached PlacementGroupCleaner actor and start monitoring."""
|
||||
|
||||
core_context = ray.runtime_context.get_runtime_context()
|
||||
self._controller_actor_id = core_context.get_actor_id()
|
||||
try:
|
||||
# Launch the cleaner as a detached actor so it survives controller death
|
||||
cleaner_actor_cls = ray.remote(num_cpus=0)(PlacementGroupCleaner)
|
||||
self._cleaner = cleaner_actor_cls.options(
|
||||
lifetime="detached",
|
||||
get_if_exists=False,
|
||||
).remote(
|
||||
controller_actor_id=self._controller_actor_id,
|
||||
check_interval_s=self._check_interval_s,
|
||||
get_actor_timeout_s=self._get_actor_timeout_s,
|
||||
stop_timeout=self._stop_timeout,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"PlacementGroupCleaner launched for run_id={train_run_context.run_id}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to launch PlacementGroupCleaner: {e}. "
|
||||
"Placement groups may not be cleaned up if controller exits ungracefully."
|
||||
)
|
||||
self._cleaner = None
|
||||
return
|
||||
|
||||
self._cleaner.start_monitoring.remote()
|
||||
|
||||
def after_worker_group_start(self, worker_group: "WorkerGroup"):
|
||||
"""Register the worker group's placement group with the cleaner.
|
||||
|
||||
This is called after a worker group is successfully started.
|
||||
"""
|
||||
if not self._cleaner or not self._controller_actor_id:
|
||||
logger.warning(
|
||||
"PlacementGroupCleaner not available. "
|
||||
"Placement groups may not be cleaned up if controller exits ungracefully."
|
||||
)
|
||||
return
|
||||
worker_group_state = worker_group.get_worker_group_state()
|
||||
placement_group = worker_group_state.placement_group_handle.placement_group
|
||||
|
||||
try:
|
||||
ray.get(self._cleaner.register_placement_group.remote(placement_group))
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to register placement group with cleaner: {e}. "
|
||||
"Placement group may not be cleaned up if controller dies ungracefully."
|
||||
)
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
f"Registered placement group {placement_group.id} with PlacementGroupCleaner."
|
||||
)
|
||||
|
||||
async def before_controller_shutdown(self):
|
||||
self._stop_cleaner()
|
||||
|
||||
def before_controller_abort(self):
|
||||
self._stop_cleaner()
|
||||
|
||||
def _stop_cleaner(self):
|
||||
if not self._cleaner:
|
||||
return
|
||||
|
||||
try:
|
||||
# Stop the cleaner gracefully (it won't clean up the PG)
|
||||
ray.get(self._cleaner.stop.remote(), timeout=self._stop_timeout)
|
||||
except RayActorError:
|
||||
logger.debug(
|
||||
"PlacementGroupCleaner exited before stop completed; ignoring."
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to stop PlacementGroupCleaner gracefully.")
|
||||
finally:
|
||||
self._cleaner = None
|
||||
@@ -0,0 +1,93 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional
|
||||
|
||||
import ray
|
||||
from ray.actor import ActorHandle
|
||||
from ray.train.v2._internal.constants import (
|
||||
DEFAULT_PREEMPTION_POLL_INTERVAL_S,
|
||||
PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
|
||||
)
|
||||
from ray.train.v2._internal.execution.callback import WorkerGroupCallback
|
||||
from ray.train.v2._internal.execution.preemption import PreemptionWatcher
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroup,
|
||||
WorkerGroupContext,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PreemptionCallback(WorkerGroupCallback):
|
||||
"""Manages a :class:`PreemptionWatcher` across worker-group lifecycles.
|
||||
|
||||
Spawns a fresh watcher in :meth:`after_worker_group_start` and stops it on
|
||||
every teardown path (shutdown and abort). Each worker group gets its own
|
||||
watcher and failure-domain map, so elastic resizes and restarts never
|
||||
leak stale state.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._poll_interval_s: float = float(
|
||||
os.getenv(
|
||||
PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
|
||||
str(DEFAULT_PREEMPTION_POLL_INTERVAL_S),
|
||||
)
|
||||
)
|
||||
self._watcher: Optional[ActorHandle] = None
|
||||
|
||||
def after_worker_group_start(self, worker_group: "WorkerGroup") -> None:
|
||||
# Tear down any watcher from a previous worker group first. Worker-group
|
||||
# startup can fail after this hook without running the shutdown hook, so
|
||||
# this also prevents leaking an orphaned watcher across a reschedule.
|
||||
self._stop_watcher()
|
||||
|
||||
# These handles are captured once per worker-group start. With the
|
||||
# standard backend, any worker replacement goes through a full worker
|
||||
# group restart (this hook runs again with fresh handles), so they
|
||||
# never go stale.
|
||||
# TODO(lehui): refresh worker handles on in-place replica replacement
|
||||
# when adding preemption support for replica groups (TorchFT).
|
||||
node_to_ranks: Dict[str, List[int]] = {}
|
||||
worker_actors_by_rank: Dict[int, ActorHandle] = {}
|
||||
for w in worker_group.get_workers():
|
||||
rank = w.distributed_context.world_rank
|
||||
node_to_ranks.setdefault(w.metadata.node_id, []).append(rank)
|
||||
worker_actors_by_rank[rank] = w.actor
|
||||
|
||||
watcher_cls = ray.remote(num_cpus=0, max_restarts=-1)(PreemptionWatcher)
|
||||
self._watcher = watcher_cls.remote(
|
||||
node_to_ranks=node_to_ranks,
|
||||
poll_interval_s=self._poll_interval_s,
|
||||
worker_actors_by_rank=worker_actors_by_rank,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"PreemptionCallback: started watcher for %d node(s).",
|
||||
len(node_to_ranks),
|
||||
)
|
||||
|
||||
def before_worker_group_shutdown(self, worker_group: "WorkerGroup") -> None:
|
||||
self._stop_watcher()
|
||||
|
||||
def after_worker_group_abort(
|
||||
self, worker_group_context: "WorkerGroupContext"
|
||||
) -> None:
|
||||
# abort() doesn't run the shutdown hook, so tear the watcher down here
|
||||
# too — otherwise it keeps polling GCS until the cluster reaps it.
|
||||
self._stop_watcher()
|
||||
|
||||
def _stop_watcher(self) -> None:
|
||||
if self._watcher is None:
|
||||
return
|
||||
watcher = self._watcher
|
||||
self._watcher = None
|
||||
# Force-kill (non-blocking) rather than a synchronous graceful stop, so
|
||||
# we never block the controller's event loop. The watcher's daemon poll
|
||||
# thread dies with the actor process and holds no external resources.
|
||||
try:
|
||||
ray.kill(watcher)
|
||||
except Exception:
|
||||
logger.warning("Failed to kill PreemptionWatcher actor.", exc_info=True)
|
||||
@@ -0,0 +1,221 @@
|
||||
import importlib
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Dict, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import Dataset
|
||||
|
||||
import ray
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ControllerCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.controller.state import (
|
||||
AbortedState,
|
||||
ErroredState,
|
||||
FinishedState,
|
||||
ReschedulingState,
|
||||
ResizingState,
|
||||
RestartingState,
|
||||
RunningState,
|
||||
SchedulingState,
|
||||
ShuttingDownState,
|
||||
TrainControllerState,
|
||||
)
|
||||
from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
|
||||
ResizeDecision,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroup,
|
||||
WorkerGroupContext,
|
||||
WorkerGroupState,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group.poll import WorkerGroupPollStatus
|
||||
from ray.train.v2._internal.logging.logging import (
|
||||
get_train_application_controller_log_path,
|
||||
)
|
||||
from ray.train.v2._internal.state.state_manager import TrainStateManager
|
||||
from ray.train.v2._internal.util import TrainingFramework
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _get_framework_version(framework: Optional[TrainingFramework]):
|
||||
versions = {}
|
||||
|
||||
try:
|
||||
import ray
|
||||
|
||||
versions["ray"] = ray.__version__
|
||||
except ImportError:
|
||||
logger.warning("Failed to collect ray version on worker.")
|
||||
|
||||
if framework is None:
|
||||
return versions
|
||||
|
||||
for module_name in framework.module_names():
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
versions[module_name] = module.__version__
|
||||
except ModuleNotFoundError:
|
||||
# Module is not installed, skip without recording a version.
|
||||
continue
|
||||
except Exception:
|
||||
logger.warning(f"Failed to collect {module_name} version on worker.")
|
||||
continue
|
||||
|
||||
return versions
|
||||
|
||||
|
||||
class StateManagerCallback(ControllerCallback, WorkerGroupCallback):
|
||||
def __init__(self, datasets: Dict[str, "Dataset"]):
|
||||
self._datasets = datasets
|
||||
|
||||
def after_controller_start(self, train_run_context: TrainRunContext):
|
||||
self._state_manager = TrainStateManager()
|
||||
self._run_name = train_run_context.get_run_config().name
|
||||
self._run_id = train_run_context.run_id
|
||||
|
||||
# TODO: Should this be generated by the caller?
|
||||
# NOTE: These must be called on the Controller.
|
||||
# The Callback is first initialized on the Driver.
|
||||
core_context = ray.runtime_context.get_runtime_context()
|
||||
self._job_id = core_context.get_job_id()
|
||||
self._controller_actor_id = core_context.get_actor_id()
|
||||
controller_log_file_path = get_train_application_controller_log_path()
|
||||
self._state_manager.create_train_run(
|
||||
id=self._run_id,
|
||||
name=self._run_name,
|
||||
job_id=self._job_id,
|
||||
controller_actor_id=self._controller_actor_id,
|
||||
controller_log_file_path=controller_log_file_path,
|
||||
run_config=train_run_context.run_config,
|
||||
train_loop_config=train_run_context.train_loop_config,
|
||||
scaling_config=train_run_context.scaling_config,
|
||||
backend_config=train_run_context.backend_config,
|
||||
datasets=self._datasets,
|
||||
dataset_config=train_run_context.dataset_config,
|
||||
)
|
||||
|
||||
def after_controller_state_update(
|
||||
self,
|
||||
previous_state: TrainControllerState,
|
||||
current_state: TrainControllerState,
|
||||
):
|
||||
if previous_state._state_type == current_state._state_type:
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"[State Transition] {previous_state._state_type.state_name} -> "
|
||||
f"{current_state._state_type.state_name}."
|
||||
)
|
||||
|
||||
if isinstance(current_state, SchedulingState):
|
||||
# TODO: This should probably always be ResizeDecision.
|
||||
if isinstance(current_state.scaling_decision, ResizeDecision):
|
||||
resize_decision = current_state.scaling_decision
|
||||
else:
|
||||
resize_decision = None
|
||||
|
||||
self._state_manager.update_train_run_scheduling(
|
||||
run_id=self._run_id,
|
||||
resize_decision=resize_decision,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, RunningState):
|
||||
self._state_manager.update_train_run_running(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, RestartingState):
|
||||
self._state_manager.update_train_run_restarting(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, ResizingState):
|
||||
self._state_manager.update_train_run_resizing(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, ErroredState):
|
||||
self._state_manager.update_train_run_errored(
|
||||
run_id=self._run_id,
|
||||
status_detail=str(current_state.training_failed_error),
|
||||
)
|
||||
|
||||
elif isinstance(current_state, FinishedState):
|
||||
self._state_manager.update_train_run_finished(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, AbortedState):
|
||||
self._state_manager.update_train_run_aborted(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, ReschedulingState):
|
||||
# substate of SchedulingState
|
||||
pass
|
||||
|
||||
elif isinstance(current_state, ShuttingDownState):
|
||||
# substate of RunningState
|
||||
pass
|
||||
|
||||
def before_worker_group_start(self, worker_group_context: WorkerGroupContext):
|
||||
self._state_manager.create_train_run_attempt(
|
||||
run_id=self._run_id,
|
||||
attempt_id=worker_group_context.run_attempt_id,
|
||||
num_workers=worker_group_context.num_workers,
|
||||
resources_per_worker=worker_group_context.resources_per_worker,
|
||||
)
|
||||
|
||||
def after_worker_group_start(self, worker_group: WorkerGroup):
|
||||
worker_group_context: WorkerGroupContext = (
|
||||
worker_group.get_worker_group_context()
|
||||
)
|
||||
worker_group_state: WorkerGroupState = worker_group.get_worker_group_state()
|
||||
self._state_manager.update_train_run_attempt_running(
|
||||
run_id=self._run_id,
|
||||
attempt_id=worker_group_context.run_attempt_id,
|
||||
workers=worker_group_state.workers,
|
||||
)
|
||||
|
||||
# Update train run framework version
|
||||
framework = self._state_manager.get_train_run_framework(self._run_id)
|
||||
framework_versions = worker_group.execute_single(
|
||||
0, _get_framework_version, framework
|
||||
)
|
||||
self._state_manager.update_train_run_framework_versions(
|
||||
run_id=self._run_id,
|
||||
framework_versions=framework_versions,
|
||||
)
|
||||
|
||||
def before_worker_group_shutdown(self, worker_group: WorkerGroup):
|
||||
worker_group_context: WorkerGroupContext = (
|
||||
worker_group.get_worker_group_context()
|
||||
)
|
||||
# TODO: Consider passing error reason directly to the callback.
|
||||
# Something along the lines of:
|
||||
# WorkerGroup.shutdown(reason)
|
||||
# -> WorkerGroupCallback.before_worker_group_shutdown(reason)
|
||||
worker_group_poll_status: Optional[
|
||||
WorkerGroupPollStatus
|
||||
] = worker_group.get_latest_poll_status()
|
||||
if worker_group_poll_status and worker_group_poll_status.errors:
|
||||
self._state_manager.update_train_run_attempt_errored(
|
||||
run_id=self._run_id,
|
||||
attempt_id=worker_group_context.run_attempt_id,
|
||||
status_detail=worker_group_poll_status.get_error_string(),
|
||||
)
|
||||
else:
|
||||
self._state_manager.update_train_run_attempt_finished(
|
||||
run_id=self._run_id,
|
||||
attempt_id=worker_group_context.run_attempt_id,
|
||||
)
|
||||
|
||||
def before_worker_group_abort(self, worker_group_context: WorkerGroupContext):
|
||||
self._state_manager.update_train_run_attempt_aborted(
|
||||
self._run_id,
|
||||
worker_group_context.run_attempt_id,
|
||||
)
|
||||
@@ -0,0 +1,54 @@
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ReportCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.training_report import _TrainingReport
|
||||
from ray.train.v2._internal.execution.worker_group import WorkerGroupPollStatus
|
||||
from ray.train.v2.api.callback import UserCallback
|
||||
|
||||
|
||||
class UserCallbackHandler(WorkerGroupCallback, ReportCallback):
|
||||
"""Responsible for calling methods of subscribers implementing
|
||||
the `UserCallback` interface.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, user_callbacks: List[UserCallback], train_run_context: TrainRunContext
|
||||
):
|
||||
self._user_callbacks = user_callbacks
|
||||
self._train_run_context = train_run_context
|
||||
|
||||
# --------------------------
|
||||
# ReportCallback
|
||||
# --------------------------
|
||||
|
||||
def after_report(
|
||||
self,
|
||||
training_report: _TrainingReport,
|
||||
metrics: List[Dict[str, Any]],
|
||||
):
|
||||
for user_callback in self._user_callbacks:
|
||||
user_callback.after_report(
|
||||
run_context=self._train_run_context,
|
||||
metrics=metrics,
|
||||
checkpoint=training_report.checkpoint,
|
||||
)
|
||||
|
||||
# --------------------------
|
||||
# WorkerGroupCallback
|
||||
# --------------------------
|
||||
|
||||
def after_worker_group_poll_status(
|
||||
self, worker_group_status: WorkerGroupPollStatus
|
||||
):
|
||||
if not worker_group_status.errors:
|
||||
return
|
||||
|
||||
for user_callback in self._user_callbacks:
|
||||
user_callback.after_exception(
|
||||
run_context=self._train_run_context,
|
||||
worker_exceptions=worker_group_status.errors,
|
||||
)
|
||||
@@ -0,0 +1,31 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ReplicaGroupCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import get_train_context
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
ExecutionGroup,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WorkingDirectorySetupCallback(ReplicaGroupCallback, WorkerGroupCallback):
|
||||
def after_execution_group_start(self, execution_group: ExecutionGroup):
|
||||
"""Shared logic for setting up the working directory on an execution group."""
|
||||
|
||||
def chdir_to_working_dir() -> None:
|
||||
"""Create the local working directory for the experiment."""
|
||||
local_working_directory = (
|
||||
get_train_context().get_storage().local_working_directory
|
||||
)
|
||||
os.makedirs(local_working_directory, exist_ok=True)
|
||||
logger.debug(
|
||||
f"Changing the working directory to: {local_working_directory}"
|
||||
)
|
||||
os.chdir(local_working_directory)
|
||||
|
||||
execution_group.execute(chdir_to_working_dir)
|
||||
Reference in New Issue
Block a user