chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,527 @@
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union
|
||||
|
||||
import pyarrow.fs
|
||||
|
||||
from ray.air.config import (
|
||||
FailureConfig as FailureConfigV1,
|
||||
ScalingConfig as ScalingConfigV1,
|
||||
)
|
||||
from ray.runtime_env import RuntimeEnv
|
||||
from ray.train.v2._internal.constants import _DEPRECATED
|
||||
from ray.train.v2._internal.execution.storage import StorageContext
|
||||
from ray.train.v2._internal.migration_utils import (
|
||||
FAIL_FAST_DEPRECATION_MESSAGE,
|
||||
TRAINER_RESOURCES_DEPRECATION_MESSAGE,
|
||||
)
|
||||
from ray.train.v2._internal.util import date_str
|
||||
from ray.util.annotations import PublicAPI
|
||||
from ray.util.tpu import get_tpu_worker_resources
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train import UserCallback
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScalingConfig(ScalingConfigV1):
|
||||
"""Configuration for scaling training.
|
||||
|
||||
Args:
|
||||
num_workers: The number of workers (Ray actors) to launch.
|
||||
Each worker will reserve 1 CPU by default. The number of CPUs
|
||||
reserved by each worker can be overridden with the
|
||||
``resources_per_worker`` argument. If the number of workers is 0,
|
||||
the training function will run in local mode, meaning the training
|
||||
function runs in the same process. To enable elasticity, provide a
|
||||
``(min_workers, max_workers)`` tuple of ints.
|
||||
elastic_resize_monitor_interval_s: While the worker group is healthy,
|
||||
consider resizing the worker group every
|
||||
``elastic_resize_monitor_interval_s`` seconds.
|
||||
use_gpu: If True, training will be done on GPUs (1 per worker).
|
||||
Defaults to False. The number of GPUs reserved by each
|
||||
worker can be overridden with the ``resources_per_worker``
|
||||
argument.
|
||||
resources_per_worker: If specified, the resources
|
||||
defined in this Dict is reserved for each worker.
|
||||
Define the ``"CPU"`` and ``"GPU"`` keys (case-sensitive) to
|
||||
override the number of CPU or GPUs used by each worker.
|
||||
|
||||
Accepts the same resource keys that Ray uses for scheduling tasks
|
||||
and actors (see :ref:`Resources <core-resources>`):
|
||||
|
||||
- ``"CPU"``: number of logical CPUs per worker.
|
||||
- ``"GPU"``: number of logical GPUs per worker. Prefer setting
|
||||
``use_gpu=True`` (which reserves 1 GPU per worker) and only
|
||||
override this key when you need a different per-worker count.
|
||||
- ``"TPU"``: number of logical TPUs per worker, when ``use_tpu=True``.
|
||||
- ``"memory"``: heap memory reserved per worker, in bytes
|
||||
(for example, ``"memory": 1e9`` reserves 1 GB per worker).
|
||||
- Any :ref:`custom resource <custom-resources>` name configured on
|
||||
your cluster (for example, ``"special_hardware": 1``).
|
||||
|
||||
Keys are case-sensitive: use ``"CPU"``, ``"GPU"``, and ``"TPU"``
|
||||
(uppercase), and ``"memory"`` (lowercase).
|
||||
placement_strategy: The placement strategy to use for the
|
||||
placement group of the Ray actors. See :ref:`Placement Group
|
||||
Strategies <pgroup-strategy>` for the possible options.
|
||||
label_selector: A list of label selectors for Ray Train worker placement.
|
||||
If a single label selector is provided, it will be applied to all Ray Train workers.
|
||||
If a list is provided, it must be the same length as the max number of Ray Train workers.
|
||||
accelerator_type: [Experimental] If specified, Ray Train will launch the
|
||||
training coordinator and workers on the nodes with the specified type
|
||||
of accelerators.
|
||||
See :ref:`the available accelerator types <accelerator_types>`.
|
||||
Ensure that your cluster has instances with the specified accelerator type
|
||||
or is able to autoscale to fulfill the request. This field is required
|
||||
when `use_tpu` is True and `num_workers` is greater than 1.
|
||||
use_tpu: [Experimental] If True, training will be done on TPUs (1 TPU VM
|
||||
per worker). Defaults to False. The number of TPUs reserved by each
|
||||
worker can be overridden with the ``resources_per_worker``
|
||||
argument. This arg enables SPMD execution of the training workload.
|
||||
topology: [Experimental] If specified, Ray Train will launch the training
|
||||
coordinator and workers on nodes with the specified topology. Topology is
|
||||
auto-detected for TPUs and added as Ray node labels. This arg enables
|
||||
SPMD execution of the training workload. This field is required
|
||||
when `use_tpu` is True and `num_workers` is greater than 1.
|
||||
"""
|
||||
|
||||
num_workers: Union[int, Tuple[int, int]] = 1
|
||||
trainer_resources: Optional[dict] = None
|
||||
label_selector: Optional[Union[Dict[str, str], List[Dict[str, str]]]] = None
|
||||
|
||||
# Accelerator specific fields.
|
||||
use_tpu: Union[bool] = False
|
||||
topology: Optional[str] = None
|
||||
|
||||
# Elasticity specific fields.
|
||||
elastic_resize_monitor_interval_s: float = 60.0
|
||||
|
||||
def __post_init__(self):
|
||||
if self.trainer_resources is not None:
|
||||
raise DeprecationWarning(TRAINER_RESOURCES_DEPRECATION_MESSAGE)
|
||||
|
||||
is_fixed = isinstance(self.num_workers, int)
|
||||
is_elastic = (
|
||||
isinstance(self.num_workers, tuple)
|
||||
and len(self.num_workers) == 2
|
||||
and all(isinstance(x, int) for x in self.num_workers)
|
||||
)
|
||||
if not (is_fixed or is_elastic):
|
||||
raise ValueError(
|
||||
"ScalingConfig(num_workers) must be an int or a tuple of two ints."
|
||||
)
|
||||
if self.elastic_resize_monitor_interval_s < 0:
|
||||
raise ValueError(
|
||||
"ScalingConfig(elastic_resize_monitor_interval_s) must be non-negative."
|
||||
)
|
||||
if self.min_workers < 0:
|
||||
raise ValueError(
|
||||
f"Invalid ScalingConfig(num_workers={self.num_workers}): "
|
||||
"Number of workers cannot be negative."
|
||||
)
|
||||
if self.min_workers > self.max_workers:
|
||||
raise ValueError(
|
||||
f"Invalid ScalingConfig(num_workers={self.num_workers}): "
|
||||
f"min_workers={self.min_workers} must be <= max_workers={self.max_workers}."
|
||||
)
|
||||
|
||||
self._validate_tpu_config()
|
||||
|
||||
if (
|
||||
isinstance(self.label_selector, list)
|
||||
and len(self.label_selector) != self.max_workers
|
||||
):
|
||||
raise ValueError(
|
||||
"If `label_selector` is a list, it must be the same length as "
|
||||
"`max_workers` (or `num_workers` when fixed)."
|
||||
)
|
||||
|
||||
if self.num_workers == 0:
|
||||
logger.info(
|
||||
"Running in local mode. The training function will run in the same process. "
|
||||
"If you are using it and running into issues please file a report at "
|
||||
"https://github.com/ray-project/ray/issues."
|
||||
)
|
||||
|
||||
super().__post_init__()
|
||||
|
||||
@property
|
||||
def elasticity_enabled(self) -> bool:
|
||||
return isinstance(self.num_workers, tuple)
|
||||
|
||||
@property
|
||||
def min_workers(self) -> int:
|
||||
return (
|
||||
self.num_workers
|
||||
if isinstance(self.num_workers, int)
|
||||
else self.num_workers[0]
|
||||
)
|
||||
|
||||
@property
|
||||
def max_workers(self) -> int:
|
||||
return (
|
||||
self.num_workers
|
||||
if isinstance(self.num_workers, int)
|
||||
else self.num_workers[1]
|
||||
)
|
||||
|
||||
def _label_selector_per_worker(
|
||||
self, num_workers: int
|
||||
) -> Optional[List[Dict[str, str]]]:
|
||||
"""Normalize ``label_selector`` into a per-worker list of length ``num_workers``.
|
||||
|
||||
- ``None`` -> ``None`` (no constraint; downstream consumers — the
|
||||
placement-group path and the autoscaling coordinator — both
|
||||
accept ``None`` and treat it as "no label requirement").
|
||||
- ``Dict`` -> the same dict replicated for each worker
|
||||
- ``List`` -> the first ``num_workers`` entries (validated to be
|
||||
``max_workers`` long in ``__post_init__``)
|
||||
"""
|
||||
if isinstance(self.label_selector, list):
|
||||
return [s.copy() for s in self.label_selector[:num_workers]]
|
||||
if isinstance(self.label_selector, dict):
|
||||
return [self.label_selector.copy() for _ in range(num_workers)]
|
||||
return None
|
||||
|
||||
@property
|
||||
def total_resources(self):
|
||||
"""Map of total resources required for training.
|
||||
|
||||
For elastic configs, this returns an upper bound based on max_workers.
|
||||
"""
|
||||
total_resource_map = dict(self._trainer_resources_not_none)
|
||||
for k, value in self._resources_per_worker_not_none.items():
|
||||
total_resource_map[k] = total_resource_map.get(k, 0.0) + (
|
||||
value * self.max_workers
|
||||
)
|
||||
return total_resource_map
|
||||
|
||||
def _validate_tpu_config(self):
|
||||
"""Validates configuration specifically for TPU usage."""
|
||||
max_workers = self.max_workers
|
||||
|
||||
if self.use_gpu and self.use_tpu:
|
||||
raise ValueError("Cannot specify both `use_gpu=True` and `use_tpu=True`.")
|
||||
|
||||
if not self.use_tpu:
|
||||
if self.num_tpus_per_worker > 0:
|
||||
raise ValueError(
|
||||
"`use_tpu` is False but `TPU` was found in "
|
||||
"`resources_per_worker`. Either set `use_tpu` to True or "
|
||||
"remove `TPU` from `resources_per_worker."
|
||||
)
|
||||
# If not using TPU, we are done validating TPU-specific logic.
|
||||
return
|
||||
|
||||
if self.num_tpus_per_worker == 0:
|
||||
raise ValueError(
|
||||
"`use_tpu` is True but `TPU` is set to 0 in "
|
||||
"`resources_per_worker`. Either set `use_tpu` to False or "
|
||||
"request a positive number of `TPU` in "
|
||||
"`resources_per_worker."
|
||||
)
|
||||
|
||||
if max_workers > 1:
|
||||
if not self.topology:
|
||||
raise ValueError(
|
||||
"`topology` must be specified in ScalingConfig when `use_tpu=True` "
|
||||
" and `num_workers` > 1."
|
||||
)
|
||||
if not self.accelerator_type:
|
||||
raise ValueError(
|
||||
"`accelerator_type` must be specified in ScalingConfig when "
|
||||
"`use_tpu=True` and `num_workers` > 1."
|
||||
)
|
||||
if self.label_selector:
|
||||
raise ValueError(
|
||||
"Cannot set `label_selector` when `use_tpu=True` because "
|
||||
"Ray Train automatically reserves a TPU slice with a predefined label."
|
||||
)
|
||||
|
||||
# Validate TPU resources when both topology and accelerator type are specified.
|
||||
if self.topology and self.accelerator_type:
|
||||
try:
|
||||
workers_per_slice, tpu_resources = get_tpu_worker_resources(
|
||||
topology=self.topology,
|
||||
accelerator_type=self.accelerator_type,
|
||||
resources_per_unit=self.resources_per_worker,
|
||||
num_slices=1,
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Could not parse TPU topology details for "
|
||||
f"type={self.accelerator_type}, "
|
||||
f"topology={self.topology}. Error: {e}"
|
||||
)
|
||||
|
||||
if workers_per_slice > 0 and max_workers % workers_per_slice != 0:
|
||||
raise ValueError(
|
||||
f"The configured `num_workers` ({self.num_workers}) must be a "
|
||||
f"multiple of {workers_per_slice} for the specified topology ({self.topology}). "
|
||||
"TPU workloads typically require symmetric resource distribution "
|
||||
"across all slices to function correctly."
|
||||
)
|
||||
if workers_per_slice > 0 and self.min_workers % workers_per_slice != 0:
|
||||
raise ValueError(
|
||||
f"The configured `min_workers` ({self.min_workers}) must be a "
|
||||
f"multiple of {workers_per_slice} for the specified topology ({self.topology}). "
|
||||
"TPU workloads typically require symmetric resource distribution "
|
||||
"across all slices to function correctly."
|
||||
)
|
||||
|
||||
if self.resources_per_worker is None:
|
||||
self.resources_per_worker = tpu_resources
|
||||
|
||||
@property
|
||||
def _resources_per_worker_not_none(self):
|
||||
if self.resources_per_worker is None:
|
||||
if self.use_tpu:
|
||||
return {"TPU": 1}
|
||||
|
||||
return super()._resources_per_worker_not_none
|
||||
|
||||
@property
|
||||
def _trainer_resources_not_none(self):
|
||||
return {}
|
||||
|
||||
@property
|
||||
def num_tpus_per_worker(self):
|
||||
"""The number of TPUs to set per worker."""
|
||||
return self._resources_per_worker_not_none.get("TPU", 0)
|
||||
|
||||
|
||||
@dataclass
|
||||
@PublicAPI(stability="stable")
|
||||
class CheckpointConfig:
|
||||
"""Configuration for checkpointing.
|
||||
|
||||
Default behavior is to persist all checkpoints reported with
|
||||
:meth:`ray.train.report` to disk. If ``num_to_keep`` is set,
|
||||
the default retention policy is to keep the most recent checkpoints.
|
||||
|
||||
Args:
|
||||
num_to_keep: The maximum number of checkpoints to keep.
|
||||
If you report more checkpoints than this, the oldest
|
||||
(or lowest-scoring, if ``checkpoint_score_attribute`` is set)
|
||||
checkpoint will be deleted.
|
||||
If this is ``None`` then all checkpoints will be kept. Must be >= 1.
|
||||
checkpoint_score_attribute: The attribute that will be used to
|
||||
score checkpoints to determine which checkpoints should be kept.
|
||||
This attribute must be a key from the metrics dictionary
|
||||
attached to the checkpoint. This attribute must have a numerical value.
|
||||
checkpoint_score_order: Either "max" or "min".
|
||||
If "max"/"min", then checkpoints with highest/lowest values of
|
||||
the ``checkpoint_score_attribute`` will be kept. Defaults to "max".
|
||||
checkpoint_frequency: [Deprecated]
|
||||
checkpoint_at_end: [Deprecated]
|
||||
"""
|
||||
|
||||
num_to_keep: Optional[int] = None
|
||||
checkpoint_score_attribute: Optional[str] = None
|
||||
checkpoint_score_order: Literal["max", "min"] = "max"
|
||||
checkpoint_frequency: Union[Optional[int], Literal[_DEPRECATED]] = _DEPRECATED
|
||||
checkpoint_at_end: Union[Optional[bool], Literal[_DEPRECATED]] = _DEPRECATED
|
||||
|
||||
def __post_init__(self):
|
||||
if self.checkpoint_frequency != _DEPRECATED:
|
||||
raise DeprecationWarning(
|
||||
"`checkpoint_frequency` is deprecated since it does not "
|
||||
"apply to user-defined training functions. "
|
||||
"Please remove this argument from your CheckpointConfig."
|
||||
)
|
||||
|
||||
if self.checkpoint_at_end != _DEPRECATED:
|
||||
raise DeprecationWarning(
|
||||
"`checkpoint_at_end` is deprecated since it does not "
|
||||
"apply to user-defined training functions. "
|
||||
"Please remove this argument from your CheckpointConfig."
|
||||
)
|
||||
|
||||
if self.num_to_keep is not None and self.num_to_keep <= 0:
|
||||
raise ValueError(
|
||||
f"Received invalid num_to_keep: {self.num_to_keep}. "
|
||||
"Must be None or an integer >= 1."
|
||||
)
|
||||
|
||||
if self.checkpoint_score_order not in ("max", "min"):
|
||||
raise ValueError(
|
||||
f"Received invalid checkpoint_score_order: {self.checkpoint_score_order}. "
|
||||
"Must be 'max' or 'min'."
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FailureConfig(FailureConfigV1):
|
||||
"""Configuration related to failure handling of each training run.
|
||||
|
||||
Args:
|
||||
max_failures: Tries to recover a run from training worker errors at least this many times.
|
||||
Will recover from the latest checkpoint if present.
|
||||
Setting to -1 will lead to infinite recovery retries.
|
||||
Setting to 0 will disable retries. Defaults to 0.
|
||||
controller_failure_limit: [DeveloperAPI] The maximum number of controller failures to tolerate.
|
||||
Setting to -1 will lead to infinite controller retries.
|
||||
Setting to 0 will disable controller retries. Defaults to -1.
|
||||
"""
|
||||
|
||||
fail_fast: Union[bool, str] = _DEPRECATED
|
||||
controller_failure_limit: int = -1
|
||||
|
||||
def __post_init__(self):
|
||||
if self.fail_fast != _DEPRECATED:
|
||||
raise DeprecationWarning(FAIL_FAST_DEPRECATION_MESSAGE)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
@dataclass
|
||||
class LoggingConfig:
|
||||
"""Configuration for Ray Train's logging behavior.
|
||||
|
||||
Args:
|
||||
log_level: The log level for Ray Train's internal ``ray.train`` logs
|
||||
on console output and application-level log files. Accepts standard
|
||||
Python logging level names. Defaults to ``"INFO"``.
|
||||
System-level log files always capture all levels (DEBUG and above),
|
||||
and the ``ray`` logger (set by ``ray.init()``) and root logger
|
||||
are unaffected.
|
||||
"""
|
||||
|
||||
log_level: str = "INFO"
|
||||
|
||||
def __post_init__(self):
|
||||
valid_levels = set(logging._nameToLevel)
|
||||
if (
|
||||
not isinstance(self.log_level, str)
|
||||
or self.log_level.upper() not in valid_levels
|
||||
):
|
||||
raise ValueError(
|
||||
f"Invalid log_level: {self.log_level!r}. "
|
||||
f"Must be one of: {', '.join(repr(x) for x in sorted(valid_levels))}."
|
||||
)
|
||||
self.log_level = self.log_level.upper()
|
||||
|
||||
|
||||
@dataclass
|
||||
@PublicAPI(stability="stable")
|
||||
class RunConfig:
|
||||
"""Runtime configuration for training runs.
|
||||
|
||||
Args:
|
||||
name: Name of the trial or experiment. If not provided, will be deduced
|
||||
from the Trainable.
|
||||
storage_path: Path where all results and checkpoints are persisted.
|
||||
Can be a local directory or a destination on cloud storage.
|
||||
For multi-node training/tuning runs, this must be set to a
|
||||
shared storage location (e.g., S3, NFS).
|
||||
This defaults to the local ``~/ray_results`` directory.
|
||||
storage_filesystem: A custom filesystem to use for storage.
|
||||
If this is provided, `storage_path` should be a path with its
|
||||
prefix stripped (e.g., `s3://bucket/path` -> `bucket/path`).
|
||||
failure_config: Failure mode configuration.
|
||||
checkpoint_config: Checkpointing configuration.
|
||||
callbacks: [DeveloperAPI] A list of callbacks that the Ray Train controller
|
||||
will invoke during training.
|
||||
worker_runtime_env: [DeveloperAPI] Runtime environment configuration
|
||||
for all Ray Train worker actors.
|
||||
logging_config: Configuration for Ray Train's logging behavior.
|
||||
See :class:`LoggingConfig` for details.
|
||||
"""
|
||||
|
||||
name: Optional[str] = None
|
||||
storage_path: Optional[str] = None
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None
|
||||
failure_config: Optional[FailureConfig] = None
|
||||
checkpoint_config: Optional[CheckpointConfig] = None
|
||||
callbacks: Optional[List["UserCallback"]] = None
|
||||
worker_runtime_env: Optional[Union[dict, RuntimeEnv]] = None
|
||||
logging_config: Optional[LoggingConfig] = None
|
||||
|
||||
sync_config: str = _DEPRECATED
|
||||
verbose: str = _DEPRECATED
|
||||
stop: str = _DEPRECATED
|
||||
progress_reporter: str = _DEPRECATED
|
||||
log_to_file: str = _DEPRECATED
|
||||
|
||||
def __post_init__(self):
|
||||
from ray.train.constants import DEFAULT_STORAGE_PATH
|
||||
|
||||
if self.storage_path is None:
|
||||
self.storage_path = DEFAULT_STORAGE_PATH
|
||||
|
||||
if not self.failure_config:
|
||||
self.failure_config = FailureConfig()
|
||||
|
||||
if not self.checkpoint_config:
|
||||
self.checkpoint_config = CheckpointConfig()
|
||||
|
||||
if not self.logging_config:
|
||||
self.logging_config = LoggingConfig()
|
||||
|
||||
if isinstance(self.storage_path, Path):
|
||||
self.storage_path = self.storage_path.as_posix()
|
||||
|
||||
run_config_deprecation_message = (
|
||||
"`RunConfig({})` is deprecated. This configuration was a "
|
||||
"Ray Tune API that did not support Ray Train usage well, "
|
||||
"so we are dropping support going forward. "
|
||||
"If you heavily rely on these configurations, "
|
||||
"you can run Ray Train as a single Ray Tune trial. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
unsupported_params = [
|
||||
"sync_config",
|
||||
"verbose",
|
||||
"stop",
|
||||
"progress_reporter",
|
||||
"log_to_file",
|
||||
]
|
||||
for param in unsupported_params:
|
||||
if getattr(self, param) != _DEPRECATED:
|
||||
raise DeprecationWarning(run_config_deprecation_message.format(param))
|
||||
|
||||
if not self.name:
|
||||
self.name = f"ray_train_run-{date_str()}"
|
||||
|
||||
self.callbacks = self.callbacks or []
|
||||
self.worker_runtime_env = self.worker_runtime_env or {}
|
||||
|
||||
from ray.train.v2.api.callback import RayTrainCallback
|
||||
|
||||
if not all(isinstance(cb, RayTrainCallback) for cb in self.callbacks):
|
||||
raise ValueError(
|
||||
"All callbacks must be instances of `ray.train.UserCallback`. "
|
||||
"Passing in a Ray Tune callback is no longer supported. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
if not isinstance(self.checkpoint_config, CheckpointConfig):
|
||||
raise ValueError(
|
||||
f"Invalid `CheckpointConfig` type: {self.checkpoint_config.__class__}. "
|
||||
"Use `ray.train.CheckpointConfig` instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
if not isinstance(self.failure_config, FailureConfig):
|
||||
raise ValueError(
|
||||
f"Invalid `FailureConfig` type: {self.failure_config.__class__}. "
|
||||
"Use `ray.train.FailureConfig` instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
@cached_property
|
||||
def storage_context(self) -> StorageContext:
|
||||
return StorageContext(
|
||||
storage_path=self.storage_path,
|
||||
experiment_dir_name=self.name,
|
||||
storage_filesystem=self.storage_filesystem,
|
||||
)
|
||||
Reference in New Issue
Block a user