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 `): - ``"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 ` 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 ` 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 `. 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, )