chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,53 @@
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class RayTrainCallback:
|
||||
"""Base Ray Train callback interface."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class UserCallback(RayTrainCallback):
|
||||
"""Callback interface for custom user-defined callbacks to handling events
|
||||
during training.
|
||||
|
||||
This callback is called on the Ray Train controller process, not on the
|
||||
worker processes.
|
||||
"""
|
||||
|
||||
def after_report(
|
||||
self,
|
||||
run_context: TrainRunContext,
|
||||
metrics: List[Dict[str, Any]],
|
||||
checkpoint: Optional[Checkpoint],
|
||||
):
|
||||
"""Called after all workers have reported a metric + checkpoint
|
||||
via `ray.train.report`.
|
||||
|
||||
Args:
|
||||
run_context: The `TrainRunContext` for the current training run.
|
||||
metrics: A list of metric dictionaries reported by workers,
|
||||
where metrics[i] is the metrics dict reported by worker i.
|
||||
checkpoint: A Checkpoint object that has been persisted to
|
||||
storage. This is None if no workers reported a checkpoint
|
||||
(e.g. `ray.train.report(metrics, checkpoint=None)`).
|
||||
"""
|
||||
pass
|
||||
|
||||
def after_exception(
|
||||
self, run_context: TrainRunContext, worker_exceptions: Dict[int, Exception]
|
||||
):
|
||||
"""Called after one or more workers have raised an exception.
|
||||
|
||||
Args:
|
||||
run_context: The `TrainRunContext` for the current training run.
|
||||
worker_exceptions: A dict from worker world rank to the exception
|
||||
raised by that worker.
|
||||
"""
|
||||
pass
|
||||
@@ -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,
|
||||
)
|
||||
@@ -0,0 +1,261 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict
|
||||
|
||||
from ray.train.v2._internal.execution.context import (
|
||||
get_train_context as get_internal_train_context,
|
||||
)
|
||||
from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
class TrainContext(ABC):
|
||||
"""Abstract interface for training context."""
|
||||
|
||||
@Deprecated
|
||||
def get_metadata(self) -> Dict[str, Any]:
|
||||
"""[Deprecated] User metadata dict passed to the Trainer constructor."""
|
||||
from ray.train.context import _GET_METADATA_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(_GET_METADATA_DEPRECATION_MESSAGE)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_name(self) -> str:
|
||||
"""[Deprecated] Trial name for the corresponding trial."""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_name")
|
||||
)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_id(self) -> str:
|
||||
"""[Deprecated] Trial id for the corresponding trial."""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_id")
|
||||
)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_resources(self):
|
||||
"""[Deprecated] Trial resources for the corresponding trial."""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_resources")
|
||||
)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_dir(self) -> str:
|
||||
"""[Deprecated] Log directory corresponding to the trial directory for a Tune session.
|
||||
This is deprecated for Ray Train and should no longer be called in Ray Train workers.
|
||||
|
||||
If this directory is needed, please pass it into the `train_loop_config` directly.
|
||||
"""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_dir")
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def get_experiment_name(self) -> str:
|
||||
"""Experiment name for the corresponding trial."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_world_size(self) -> int:
|
||||
"""Get the current world size (i.e. total number of workers) for this run.
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
NUM_WORKERS = 2
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
assert ray.train.get_context().get_world_size() == NUM_WORKERS
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=NUM_WORKERS),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_world_rank(self) -> int:
|
||||
"""Get the world rank of this worker.
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
print("Worker 0")
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_local_rank(self) -> int:
|
||||
"""Get the local rank of this worker (rank of the worker on its node).
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
if ray.train.get_context().get_local_rank() == 0:
|
||||
print("Local rank 0 worker")
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_local_world_size(self) -> int:
|
||||
"""Get the local world size of this node (i.e. number of workers on this node).
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker():
|
||||
print(ray.train.get_context().get_local_world_size())
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Returns:
|
||||
The number of workers running on this node.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_node_rank(self) -> int:
|
||||
"""Get the rank of this node.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker():
|
||||
print(ray.train.get_context().get_node_rank())
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=1),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Returns:
|
||||
The rank of this node among the nodes participating in training.
|
||||
"""
|
||||
pass
|
||||
|
||||
@DeveloperAPI
|
||||
@abstractmethod
|
||||
def get_storage(self):
|
||||
"""Returns the :class:`~ray.train._internal.storage.StorageContext` storage
|
||||
context which gives advanced access to the filesystem and paths
|
||||
configured through `RunConfig`.
|
||||
|
||||
NOTE: This is a DeveloperAPI, and the `StorageContext` interface may change
|
||||
without notice between minor versions.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DistributedTrainContext(TrainContext):
|
||||
"""Implementation of TrainContext for distributed mode."""
|
||||
|
||||
def get_experiment_name(self) -> str:
|
||||
return get_internal_train_context().get_experiment_name()
|
||||
|
||||
def get_world_size(self) -> int:
|
||||
return get_internal_train_context().get_world_size()
|
||||
|
||||
def get_world_rank(self) -> int:
|
||||
return get_internal_train_context().get_world_rank()
|
||||
|
||||
def get_local_rank(self) -> int:
|
||||
return get_internal_train_context().get_local_rank()
|
||||
|
||||
def get_local_world_size(self) -> int:
|
||||
return get_internal_train_context().get_local_world_size()
|
||||
|
||||
def get_node_rank(self) -> int:
|
||||
return get_internal_train_context().get_node_rank()
|
||||
|
||||
def get_storage(self):
|
||||
return get_internal_train_context().get_storage()
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class LocalTrainContext(TrainContext):
|
||||
"""Implementation of TrainContext for local mode."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
experiment_name: str,
|
||||
world_size: int = 1,
|
||||
world_rank: int = 0,
|
||||
local_rank: int = 0,
|
||||
local_world_size: int = 1,
|
||||
node_rank: int = 0,
|
||||
):
|
||||
self.experiment_name = experiment_name
|
||||
self.world_size = world_size
|
||||
self.world_rank = world_rank
|
||||
self.local_rank = local_rank
|
||||
self.local_world_size = local_world_size
|
||||
self.node_rank = node_rank
|
||||
|
||||
def get_experiment_name(self) -> str:
|
||||
return self.experiment_name
|
||||
|
||||
def get_world_size(self) -> int:
|
||||
return self.world_size
|
||||
|
||||
def get_world_rank(self) -> int:
|
||||
return self.world_rank
|
||||
|
||||
def get_local_rank(self) -> int:
|
||||
return self.local_rank
|
||||
|
||||
def get_local_world_size(self) -> int:
|
||||
return self.local_world_size
|
||||
|
||||
def get_node_rank(self) -> int:
|
||||
return self.node_rank
|
||||
|
||||
def get_storage(self):
|
||||
raise NotImplementedError("Local storage context not yet implemented. ")
|
||||
@@ -0,0 +1,335 @@
|
||||
import logging
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import ray
|
||||
from ray._common.constants import RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR
|
||||
from ray._common.usage import usage_lib
|
||||
from ray._private.ray_constants import env_bool
|
||||
from ray.actor import ActorHandle
|
||||
from ray.air._internal.usage import tag_train_v2_trainer
|
||||
from ray.train import (
|
||||
BackendConfig,
|
||||
Checkpoint,
|
||||
DataConfig,
|
||||
Result,
|
||||
RunConfig,
|
||||
ScalingConfig,
|
||||
)
|
||||
from ray.train.base_trainer import (
|
||||
_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING,
|
||||
_TRAINER_RESTORE_DEPRECATION_WARNING,
|
||||
)
|
||||
from ray.train.constants import RAY_CHDIR_TO_TRIAL_DIR, RAY_TRAIN_ENABLE_STATE_TRACKING
|
||||
from ray.train.context import _GET_METADATA_DEPRECATION_MESSAGE
|
||||
from ray.train.v2._internal.callbacks import (
|
||||
AcceleratorSetupCallback,
|
||||
BackendSetupCallback,
|
||||
DatasetsCallback,
|
||||
WorkingDirectorySetupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.callbacks.env_callback import _initialize_env_callbacks
|
||||
from ray.train.v2._internal.callbacks.metrics import (
|
||||
ControllerMetricsCallback,
|
||||
WorkerMetricsCallback,
|
||||
)
|
||||
from ray.train.v2._internal.callbacks.placement_group_callback import (
|
||||
PlacementGroupCleanerCallback,
|
||||
)
|
||||
from ray.train.v2._internal.callbacks.state_manager import StateManagerCallback
|
||||
from ray.train.v2._internal.callbacks.user_callback import UserCallbackHandler
|
||||
from ray.train.v2._internal.constants import (
|
||||
DEFAULT_RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_VALUE,
|
||||
METRICS_ENABLED_ENV_VAR,
|
||||
V2_ENABLED_ENV_VAR,
|
||||
get_env_vars_to_propagate,
|
||||
is_v2_enabled,
|
||||
)
|
||||
from ray.train.v2._internal.data_integration.interfaces import GenDataset
|
||||
from ray.train.v2._internal.execution.callback import RayTrainCallback
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.controller import TrainController
|
||||
from ray.train.v2._internal.execution.failure_handling import create_failure_policy
|
||||
from ray.train.v2._internal.execution.local_mode.utils import LocalController
|
||||
from ray.train.v2._internal.execution.scaling_policy import create_scaling_policy
|
||||
from ray.train.v2._internal.util import ObjectRefWrapper, construct_train_func
|
||||
from ray.train.v2.api.callback import UserCallback
|
||||
from ray.train.v2.api.validation_config import ValidationConfig
|
||||
from ray.util.annotations import Deprecated, DeveloperAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DataParallelTrainer:
|
||||
"""Base class for distributed data parallel training on Ray.
|
||||
|
||||
This class supports the SPMD parallelization pattern, where a single
|
||||
training function is executed in parallel across multiple workers,
|
||||
and different shards of data are processed by each worker.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], Any], Callable[[Dict], Any]],
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
backend_config: Optional[BackendConfig] = None,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
dataset_config: Optional[DataConfig] = None,
|
||||
# TODO: [Deprecated] Remove in future release
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
validation_config: Optional[ValidationConfig] = None,
|
||||
):
|
||||
self.run_config = run_config or RunConfig()
|
||||
self.train_loop_per_worker = train_loop_per_worker
|
||||
self.validation_config = validation_config
|
||||
self.train_loop_config = train_loop_config
|
||||
self.scaling_config = scaling_config or ScalingConfig()
|
||||
self.backend_config = backend_config or BackendConfig()
|
||||
self.datasets = datasets or {}
|
||||
self.data_config = dataset_config or DataConfig()
|
||||
|
||||
self.running_in_local_mode = self.scaling_config.num_workers == 0
|
||||
|
||||
self.train_run_context = TrainRunContext(
|
||||
run_config=self.run_config,
|
||||
train_loop_config=self.train_loop_config,
|
||||
scaling_config=self.scaling_config,
|
||||
backend_config=self.backend_config,
|
||||
dataset_config=self.data_config,
|
||||
)
|
||||
|
||||
if resume_from_checkpoint is not None:
|
||||
raise DeprecationWarning(_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING)
|
||||
|
||||
if metadata is not None:
|
||||
raise DeprecationWarning(_GET_METADATA_DEPRECATION_MESSAGE)
|
||||
|
||||
self._validate_configs()
|
||||
|
||||
usage_lib.record_library_usage("train")
|
||||
tag_train_v2_trainer(self)
|
||||
if self.scaling_config.elasticity_enabled:
|
||||
usage_lib.record_extra_usage_tag(
|
||||
usage_lib.TagKey.TRAIN_ELASTICITY_ENABLED, "1"
|
||||
)
|
||||
|
||||
def _validate_configs(self):
|
||||
if not is_v2_enabled():
|
||||
raise ValueError(
|
||||
f"Ray Train V2 must be enabled with `{V2_ENABLED_ENV_VAR}=1` "
|
||||
"when using this V2 Trainer API."
|
||||
)
|
||||
|
||||
from ray.train.v2.api.config import (
|
||||
RunConfig as RunConfigV2,
|
||||
ScalingConfig as ScalingConfigV2,
|
||||
)
|
||||
|
||||
if not isinstance(self.run_config, RunConfigV2):
|
||||
raise ValueError(
|
||||
f"Invalid `RunConfig` type: {self.run_config.__class__}. "
|
||||
"Use `ray.train.RunConfig` instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
if not isinstance(self.scaling_config, ScalingConfigV2):
|
||||
raise ValueError(
|
||||
f"Invalid `ScalingConfig` type: {self.scaling_config.__class__}. "
|
||||
"Use `ray.train.ScalingConfig` instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
def _get_train_func(self) -> Callable[[], Any]:
|
||||
return construct_train_func(
|
||||
self.train_loop_per_worker,
|
||||
config=self.train_loop_config,
|
||||
train_func_context=self.backend_config.train_func_context,
|
||||
fn_arg_name="train_loop_per_worker",
|
||||
)
|
||||
|
||||
def fit(self) -> Result:
|
||||
"""Launches the Ray Train controller to run training on workers.
|
||||
|
||||
Returns:
|
||||
A Result object containing the training result.
|
||||
|
||||
Raises:
|
||||
ray.train.TrainingFailedError: This is a union of the ControllerError and WorkerGroupError.
|
||||
This returns a :class:`ray.train.ControllerError` if internal Ray Train controller logic
|
||||
encounters a non-retryable error or reaches the controller failure limit configured in `FailureConfig`.
|
||||
This returns a :class:`ray.train.WorkerGroupError` if one or more workers fail during
|
||||
training and reaches the worker group failure limit configured in `FailureConfig(max_failures)`.
|
||||
"""
|
||||
train_fn = self._get_train_func()
|
||||
if self.running_in_local_mode:
|
||||
return self._initialize_and_run_local_controller(train_fn)
|
||||
else:
|
||||
train_fn_ref = ObjectRefWrapper(train_fn)
|
||||
|
||||
result = self._initialize_and_run_controller(
|
||||
train_fn_ref=train_fn_ref,
|
||||
scaling_policy=create_scaling_policy(self.scaling_config),
|
||||
failure_policy=create_failure_policy(self.run_config.failure_config),
|
||||
train_run_context=self.train_run_context,
|
||||
callbacks=self._create_default_callbacks(),
|
||||
validation_config=self.validation_config,
|
||||
)
|
||||
|
||||
if result.error:
|
||||
# NOTE: If the training run errored out, raise an error back to the
|
||||
# user's driver script.
|
||||
# For example, if the Train `FailurePolicy` runs out of retries,
|
||||
# and one of the workers errors. The controller will exit, and
|
||||
# the error will be raised here.
|
||||
raise result.error
|
||||
|
||||
return result
|
||||
|
||||
def _get_local_controller(self) -> LocalController:
|
||||
return LocalController(
|
||||
experiment_name=self.run_config.name,
|
||||
datasets=self.datasets,
|
||||
)
|
||||
|
||||
def _create_default_callbacks(self) -> List[RayTrainCallback]:
|
||||
# Initialize callbacks from environment variable
|
||||
callbacks = _initialize_env_callbacks()
|
||||
|
||||
accelerator_setup_callback = AcceleratorSetupCallback(
|
||||
self.backend_config, self.scaling_config
|
||||
)
|
||||
backend_setup_callback = BackendSetupCallback(self.backend_config)
|
||||
datasets_callback = DatasetsCallback(
|
||||
train_run_context=self.train_run_context,
|
||||
datasets=self.datasets,
|
||||
)
|
||||
placement_group_cleaner_callback = PlacementGroupCleanerCallback()
|
||||
callbacks.extend(
|
||||
[
|
||||
accelerator_setup_callback,
|
||||
backend_setup_callback,
|
||||
placement_group_cleaner_callback,
|
||||
datasets_callback,
|
||||
]
|
||||
)
|
||||
if env_bool(RAY_CHDIR_TO_TRIAL_DIR, True):
|
||||
working_directory_setup_callback = WorkingDirectorySetupCallback()
|
||||
callbacks.append(working_directory_setup_callback)
|
||||
|
||||
if env_bool(METRICS_ENABLED_ENV_VAR, True):
|
||||
callbacks.append(ControllerMetricsCallback())
|
||||
callbacks.append(WorkerMetricsCallback(self.train_run_context))
|
||||
|
||||
if env_bool(RAY_TRAIN_ENABLE_STATE_TRACKING, False):
|
||||
callbacks.append(StateManagerCallback(datasets=self.datasets))
|
||||
|
||||
run_config_callbacks = (
|
||||
self.run_config.callbacks if self.run_config.callbacks is not None else []
|
||||
)
|
||||
|
||||
# Add internal callback that invokes all user-defined callbacks.
|
||||
user_callbacks = [
|
||||
cb for cb in run_config_callbacks if isinstance(cb, UserCallback)
|
||||
]
|
||||
callbacks.append(
|
||||
UserCallbackHandler(
|
||||
user_callbacks=user_callbacks, train_run_context=self.train_run_context
|
||||
)
|
||||
)
|
||||
|
||||
# Append all other callbacks to the full list. This allows custom workarounds
|
||||
# built on top of internal callbacks to work.
|
||||
callbacks.extend(
|
||||
[cb for cb in run_config_callbacks if not isinstance(cb, UserCallback)]
|
||||
)
|
||||
return callbacks
|
||||
|
||||
def _initialize_and_run_local_controller(
|
||||
self, train_func: Callable[[], Any]
|
||||
) -> Result:
|
||||
return self._get_local_controller().run(train_func)
|
||||
|
||||
def _initialize_and_run_controller(self, **controller_init_kwargs) -> Result:
|
||||
env_vars = get_env_vars_to_propagate()
|
||||
env_vars.setdefault(
|
||||
RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR,
|
||||
DEFAULT_RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_VALUE,
|
||||
)
|
||||
|
||||
# Attach the controller to the node running the driver script.
|
||||
controller_actor_cls = ray.remote(
|
||||
num_cpus=0,
|
||||
label_selector={
|
||||
ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
|
||||
},
|
||||
# TODO: Extract env variables that affect controller behavior
|
||||
# and pass them as explicit args
|
||||
runtime_env={"env_vars": env_vars},
|
||||
)(TrainController)
|
||||
|
||||
controller = controller_actor_cls.remote(**controller_init_kwargs)
|
||||
|
||||
# If this is not the main thread - as is the case when running in Tune -
|
||||
# registering the SIGINT handler raises an exception.
|
||||
if threading.current_thread() is threading.main_thread():
|
||||
self._register_sigint_handler(controller)
|
||||
|
||||
ray.get(controller.run.remote())
|
||||
return ray.get(controller.get_result.remote())
|
||||
|
||||
def _register_sigint_handler(self, controller: ActorHandle[TrainController]):
|
||||
"""Register SIGINT handler so user Ctrl C gracefully aborts run."""
|
||||
sigint_count = 0
|
||||
|
||||
def sigint_handler(signum, frame):
|
||||
logger.info(
|
||||
"Received SIGINT. Gracefully aborting the training run — this "
|
||||
"may take a few seconds. To forcefully abort immediately, you "
|
||||
"can send a different signal, such as SIGKILL."
|
||||
)
|
||||
nonlocal sigint_count
|
||||
sigint_count += 1
|
||||
if sigint_count >= 3:
|
||||
logger.info(
|
||||
"Received SIGINT at least 3 times. "
|
||||
"Forcefully aborting the training run."
|
||||
)
|
||||
sys.exit(0)
|
||||
if sigint_count <= 1:
|
||||
try:
|
||||
ray.get(controller.abort.remote())
|
||||
except ray.exceptions.RayActorError:
|
||||
# We catch the error and exit 0 to indicate graceful termination.
|
||||
# However, for some reason the process still exits with 1.
|
||||
sys.exit(0)
|
||||
|
||||
signal.signal(signal.SIGINT, sigint_handler)
|
||||
|
||||
@classmethod
|
||||
@Deprecated
|
||||
def restore(cls, *args, **kwargs):
|
||||
"""[Deprecated] Restores a Train experiment from a previously
|
||||
interrupted/failed run.
|
||||
|
||||
This method is deprecated and will be removed in a future release.
|
||||
"""
|
||||
raise DeprecationWarning(_TRAINER_RESTORE_DEPRECATION_WARNING)
|
||||
|
||||
@classmethod
|
||||
@Deprecated
|
||||
def can_restore(cls, *args, **kwargs):
|
||||
"""[Deprecated] Checks if a Train experiment can be restored from
|
||||
a previously interrupted/failed run.
|
||||
|
||||
This method is deprecated and will be removed in a future release.
|
||||
"""
|
||||
raise DeprecationWarning(_TRAINER_RESTORE_DEPRECATION_WARNING)
|
||||
@@ -0,0 +1,49 @@
|
||||
from typing import Dict
|
||||
|
||||
from ray.train.v2._internal.exceptions import RayTrainError
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class TrainingFailedError(RayTrainError):
|
||||
"""Exception raised when training fails from a `trainer.fit()` call.
|
||||
This is either :class:`ray.train.WorkerGroupError` or :class:`ray.train.ControllerError`.
|
||||
"""
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class WorkerGroupError(TrainingFailedError):
|
||||
"""Exception raised from the worker group during training.
|
||||
|
||||
Args:
|
||||
error_message: A human-readable error message describing the training worker failures.
|
||||
worker_failures: A mapping from worker rank to the exception that
|
||||
occurred on that worker during training.
|
||||
"""
|
||||
|
||||
def __init__(self, error_message: str, worker_failures: Dict[int, Exception]):
|
||||
super().__init__("Training failed due to worker errors:\n" + error_message)
|
||||
self._error_message = error_message
|
||||
self.worker_failures = worker_failures
|
||||
|
||||
def __reduce__(self):
|
||||
return (self.__class__, (self._error_message, self.worker_failures))
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class ControllerError(TrainingFailedError):
|
||||
"""Exception raised when training fails due to a controller error.
|
||||
|
||||
Args:
|
||||
controller_failure: The exception that occurred on the controller.
|
||||
"""
|
||||
|
||||
def __init__(self, controller_failure: Exception):
|
||||
super().__init__(
|
||||
"Training failed due to controller error:\n" + str(controller_failure)
|
||||
)
|
||||
self.controller_failure = controller_failure
|
||||
self.with_traceback(controller_failure.__traceback__)
|
||||
|
||||
def __reduce__(self):
|
||||
return (self.__class__, (self.controller_failure,))
|
||||
@@ -0,0 +1,36 @@
|
||||
from enum import Enum
|
||||
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class CheckpointUploadMode(Enum):
|
||||
"""The manner in which we want to upload the checkpoint.
|
||||
|
||||
Members:
|
||||
ASYNC: Upload checkpoint asynchronously.
|
||||
SYNC: Upload checkpoint synchronously.
|
||||
NO_UPLOAD: Do not upload checkpoint.
|
||||
"""
|
||||
|
||||
ASYNC = "ASYNC"
|
||||
SYNC = "SYNC"
|
||||
NO_UPLOAD = "NO_UPLOAD"
|
||||
|
||||
def default_delete_local_checkpoint_after_upload(self) -> bool:
|
||||
return self == CheckpointUploadMode.ASYNC
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class CheckpointConsistencyMode(Enum):
|
||||
"""Read semantics for checkpoint retrieval during an ongoing run.
|
||||
|
||||
Members:
|
||||
COMMITTED: Block until the checkpoint from the latest ray.train.report
|
||||
has been uploaded and committed.
|
||||
VALIDATED: Block until the checkpoint from the latest ray.train.report
|
||||
has been uploaded and validated.
|
||||
"""
|
||||
|
||||
COMMITTED = "COMMITTED"
|
||||
VALIDATED = "VALIDATED"
|
||||
@@ -0,0 +1,43 @@
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Any, Dict
|
||||
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train import Checkpoint
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class ReportedCheckpointStatus(Enum):
|
||||
"""ReportedCheckpoint status enum.
|
||||
|
||||
* COMMITTED: The checkpoint is saved, and no validation was requested.
|
||||
* PENDING_VALIDATION: The checkpoint is saved, and validation is in progress.
|
||||
* VALIDATED: The checkpoint is saved, and validation is complete.
|
||||
* VALIDATION_TIMEOUT: The checkpoint is saved, and validation is timed out according to
|
||||
`ValidationTaskConfig(..., timeout_s=N)`.
|
||||
* VALIDATION_FAILED: The checkpoint is saved, and validation failed.
|
||||
"""
|
||||
|
||||
COMMITTED = "COMMITTED"
|
||||
PENDING_VALIDATION = "PENDING_VALIDATION"
|
||||
VALIDATED = "VALIDATED"
|
||||
VALIDATION_TIMEOUT = "VALIDATION_TIMEOUT"
|
||||
VALIDATION_FAILED = "VALIDATION_FAILED"
|
||||
|
||||
|
||||
@dataclass
|
||||
@PublicAPI(stability="alpha")
|
||||
class ReportedCheckpoint:
|
||||
"""A user-reported checkpoint and its associated metrics.
|
||||
|
||||
Attributes:
|
||||
checkpoint: The checkpoint reported by the user.
|
||||
metrics: The metrics associated with that checkpoint.
|
||||
status: The status of the checkpoint.
|
||||
"""
|
||||
|
||||
checkpoint: "Checkpoint"
|
||||
metrics: Dict[str, Any]
|
||||
status: ReportedCheckpointStatus
|
||||
@@ -0,0 +1,162 @@
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow
|
||||
|
||||
import ray
|
||||
from ray.air.result import Result as ResultV1
|
||||
from ray.train import Checkpoint, CheckpointConfig
|
||||
from ray.train.v2._internal.constants import CHECKPOINT_MANAGER_SNAPSHOT_FILENAME
|
||||
from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
|
||||
CheckpointManager,
|
||||
)
|
||||
from ray.train.v2._internal.execution.storage import (
|
||||
StorageContext,
|
||||
_exists_at_fs_path,
|
||||
get_fs_and_path,
|
||||
)
|
||||
from ray.train.v2.api.exceptions import TrainingFailedError
|
||||
from ray.util.annotations import Deprecated, PublicAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Result(ResultV1):
|
||||
"""The output of ``trainer.fit()``.
|
||||
|
||||
Attributes:
|
||||
metrics: The latest set of metrics reported by the training function
|
||||
via :func:`ray.train.report`.
|
||||
checkpoint: The latest checkpoint saved by the training function
|
||||
via :func:`ray.train.report`.
|
||||
return_value: The value returned by the user-defined training function on the
|
||||
rank 0 worker, or ``None`` if no value was returned or if training did
|
||||
not complete successfully. The return value must be serializable.
|
||||
metrics_dataframe: A DataFrame of metrics from all checkpoints saved
|
||||
during the run. Each row corresponds to a checkpoint.
|
||||
best_checkpoints: A list of ``(checkpoint, metrics)`` tuples for the
|
||||
best checkpoints saved during the run. The checkpoints retained
|
||||
are determined by :class:`~ray.train.CheckpointConfig`
|
||||
(by default, all checkpoints are kept).
|
||||
path: Path pointing to the run output directory on persistent storage.
|
||||
This can point to a remote storage location (e.g. S3) or to a
|
||||
local location on the head node.
|
||||
error: The execution error of the training run, if the run finished
|
||||
in error. This is a :class:`~ray.train.v2.api.exceptions.TrainingFailedError`
|
||||
wrapping the original exception.
|
||||
"""
|
||||
|
||||
checkpoint: Optional[Checkpoint]
|
||||
error: Optional[TrainingFailedError]
|
||||
best_checkpoints: Optional[List[Tuple[Checkpoint, Dict[str, Any]]]] = None
|
||||
return_value: Optional[Any] = None
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
def get_best_checkpoint(
|
||||
self, metric: str, mode: str
|
||||
) -> Optional["ray.train.Checkpoint"]:
|
||||
return super().get_best_checkpoint(metric, mode)
|
||||
|
||||
@classmethod
|
||||
def from_path(
|
||||
cls,
|
||||
path: Union[str, os.PathLike],
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
||||
) -> "Result":
|
||||
"""Restore a training result from a previously saved training run path.
|
||||
|
||||
Args:
|
||||
path: Path to the run output directory
|
||||
storage_filesystem: Optional filesystem to use for accessing the path
|
||||
|
||||
Returns:
|
||||
Result object with restored checkpoints and metrics
|
||||
"""
|
||||
fs, fs_path = get_fs_and_path(str(path), storage_filesystem)
|
||||
|
||||
# Validate that the experiment directory exists
|
||||
if not _exists_at_fs_path(fs, fs_path):
|
||||
raise RuntimeError(f"Experiment folder {fs_path} doesn't exist.")
|
||||
|
||||
# Remove trailing slashes to handle paths correctly
|
||||
# os.path.basename() returns empty string for paths with trailing slashes
|
||||
fs_path = fs_path.rstrip("/")
|
||||
storage_path, experiment_dir_name = os.path.dirname(fs_path), os.path.basename(
|
||||
fs_path
|
||||
)
|
||||
|
||||
storage_context = StorageContext(
|
||||
storage_path=storage_path,
|
||||
experiment_dir_name=experiment_dir_name,
|
||||
storage_filesystem=fs,
|
||||
read_only=True,
|
||||
)
|
||||
|
||||
# Validate that the checkpoint manager snapshot file exists
|
||||
if not _exists_at_fs_path(
|
||||
storage_context.storage_filesystem,
|
||||
storage_context.checkpoint_manager_snapshot_path,
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Failed to restore the Result object: "
|
||||
f"{CHECKPOINT_MANAGER_SNAPSHOT_FILENAME} doesn't exist in the "
|
||||
f"experiment folder. Make sure that this is an output directory created by a Ray Train run."
|
||||
)
|
||||
|
||||
checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=CheckpointConfig(),
|
||||
)
|
||||
|
||||
# When we build a Result object from checkpoints, the error is not loaded.
|
||||
return cls._from_checkpoint_manager(
|
||||
checkpoint_manager=checkpoint_manager,
|
||||
storage_context=storage_context,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _from_checkpoint_manager(
|
||||
cls,
|
||||
checkpoint_manager: CheckpointManager,
|
||||
storage_context: StorageContext,
|
||||
error: Optional[TrainingFailedError] = None,
|
||||
) -> "Result":
|
||||
"""Create a Result object from a CheckpointManager."""
|
||||
latest_checkpoint_result = checkpoint_manager.latest_checkpoint_result
|
||||
if latest_checkpoint_result:
|
||||
latest_metrics = latest_checkpoint_result.metrics
|
||||
latest_checkpoint = latest_checkpoint_result.checkpoint
|
||||
else:
|
||||
latest_metrics = None
|
||||
latest_checkpoint = None
|
||||
best_checkpoints = [
|
||||
(r.checkpoint, r.metrics)
|
||||
for r in checkpoint_manager.best_checkpoint_results
|
||||
]
|
||||
|
||||
# Provide the history of metrics attached to checkpoints as a dataframe.
|
||||
metrics_dataframe = None
|
||||
if best_checkpoints:
|
||||
metrics_dataframe = pd.DataFrame([m for _, m in best_checkpoints])
|
||||
|
||||
return Result(
|
||||
metrics=latest_metrics,
|
||||
checkpoint=latest_checkpoint,
|
||||
error=error,
|
||||
path=storage_context.experiment_fs_path,
|
||||
best_checkpoints=best_checkpoints,
|
||||
metrics_dataframe=metrics_dataframe,
|
||||
_storage_filesystem=storage_context.storage_filesystem,
|
||||
)
|
||||
|
||||
@property
|
||||
@Deprecated
|
||||
def config(self) -> Optional[Dict[str, Any]]:
|
||||
raise DeprecationWarning(
|
||||
"The `config` property for a `ray.train.Result` is deprecated, "
|
||||
"since it is only relevant in the context of Ray Tune."
|
||||
)
|
||||
@@ -0,0 +1,297 @@
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
|
||||
from ray.train.v2._internal.data_integration.interfaces import DatasetShardMetadata
|
||||
from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
|
||||
from ray.train.v2._internal.util import requires_train_worker
|
||||
from ray.train.v2.api.context import TrainContext
|
||||
from ray.train.v2.api.report_config import (
|
||||
CheckpointConsistencyMode,
|
||||
CheckpointUploadMode,
|
||||
)
|
||||
from ray.train.v2.api.validation_config import ValidationTaskConfig
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import DataIterator
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@requires_train_worker(raise_in_tune_session=True)
|
||||
def report(
|
||||
metrics: Dict[str, Any],
|
||||
checkpoint: Optional["Checkpoint"] = None,
|
||||
checkpoint_dir_name: Optional[str] = None,
|
||||
checkpoint_upload_mode: CheckpointUploadMode = CheckpointUploadMode.SYNC,
|
||||
delete_local_checkpoint_after_upload: Optional[bool] = None,
|
||||
checkpoint_upload_fn: Optional[Callable[["Checkpoint", str], "Checkpoint"]] = None,
|
||||
validation: Union[bool, ValidationTaskConfig] = False,
|
||||
):
|
||||
"""Report metrics and optionally save a checkpoint.
|
||||
|
||||
If a checkpoint is provided, it will be
|
||||
:ref:`persisted to storage <persistent-storage-guide>`.
|
||||
|
||||
If this is called in multiple distributed training workers:
|
||||
|
||||
- Only the metrics reported by the rank 0 worker will be attached to the checkpoint.
|
||||
- A checkpoint will be registered as long as one or more workers reports
|
||||
checkpoint that is not None.
|
||||
See the :ref:`checkpointing guide <train-dl-saving-checkpoints>`.
|
||||
- Checkpoints from multiple workers will be merged into one directory
|
||||
in persistent storage.
|
||||
See :ref:`the distributed checkpointing guide <train-distributed-checkpointing>`.
|
||||
|
||||
|
||||
.. warning::
|
||||
|
||||
All workers must call `ray.train.report` the same number of times
|
||||
so that Ray Train can properly synchronize the training state across
|
||||
workers. This method acts as a barrier across all workers, so be sure
|
||||
that every worker reaches this method.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import tempfile
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
def train_func(config):
|
||||
start_epoch = 0
|
||||
|
||||
for epoch in range(start_epoch, config.get("num_epochs", 10)):
|
||||
# Do training...
|
||||
|
||||
metrics = {"loss": ...}
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
# Save the checkpoint...
|
||||
# torch.save(...)
|
||||
|
||||
checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
|
||||
|
||||
# Example: Only the rank 0 worker uploads the checkpoint.
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
ray.train.report(metrics, checkpoint=checkpoint)
|
||||
else:
|
||||
ray.train.report(metrics, checkpoint=None)
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_func, scaling_config=ray.train.ScalingConfig(num_workers=2)
|
||||
)
|
||||
|
||||
Args:
|
||||
metrics: The metrics you want to report.
|
||||
checkpoint: The optional checkpoint you want to report.
|
||||
checkpoint_dir_name: Custom name for the checkpoint directory.
|
||||
If not provided, a unique directory name will be automatically generated.
|
||||
If provided, it must be unique across all checkpoints per worker to avoid
|
||||
naming collisions. Consider including identifiers such as the epoch or batch
|
||||
index in the name.
|
||||
checkpoint_upload_mode: The manner in which we want to upload the checkpoint.
|
||||
Defaults to uploading the checkpoint synchronously.
|
||||
This works when no checkpoint is provided but is not useful in that case.
|
||||
delete_local_checkpoint_after_upload: Whether to delete the checkpoint after it is uploaded.
|
||||
checkpoint_upload_fn: A user defined function that will be called with the
|
||||
checkpoint to upload it. If not provided, defaults to using the `pyarrow.fs.copy_files`
|
||||
utility for copying to the destination `storage_path`.
|
||||
validation: [Alpha] If True, triggers validation with default kwargs from validation_config.
|
||||
If a ValidationTaskConfig, validation is run using fn_kwargs merged with validation_config
|
||||
defaults, with fn_kwargs taking precedence on conflicts. If False, no validation.
|
||||
"""
|
||||
if validation and not checkpoint:
|
||||
raise ValueError("Validation requires a checkpoint to be provided.")
|
||||
|
||||
if delete_local_checkpoint_after_upload is None:
|
||||
delete_local_checkpoint_after_upload = (
|
||||
checkpoint_upload_mode.default_delete_local_checkpoint_after_upload()
|
||||
)
|
||||
|
||||
if checkpoint:
|
||||
record_extra_usage_tag(
|
||||
TagKey.TRAIN_CHECKPOINT_MODE, checkpoint_upload_mode.value
|
||||
)
|
||||
if validation:
|
||||
record_extra_usage_tag(TagKey.TRAIN_ASYNCHRONOUS_VALIDATION, "1")
|
||||
|
||||
get_train_fn_utils().report(
|
||||
metrics=metrics,
|
||||
checkpoint=checkpoint,
|
||||
checkpoint_dir_name=checkpoint_dir_name,
|
||||
checkpoint_upload_mode=checkpoint_upload_mode,
|
||||
delete_local_checkpoint_after_upload=delete_local_checkpoint_after_upload,
|
||||
checkpoint_upload_fn=checkpoint_upload_fn,
|
||||
validation=validation,
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@requires_train_worker(raise_in_tune_session=True)
|
||||
def get_context() -> TrainContext:
|
||||
"""Get or create a singleton training context.
|
||||
|
||||
The context is only available within a function passed to Ray Train.
|
||||
|
||||
See the :class:`~ray.train.TrainContext` API reference to see available methods.
|
||||
"""
|
||||
return get_train_fn_utils().get_context()
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@requires_train_worker(raise_in_tune_session=True)
|
||||
def get_checkpoint() -> Optional["Checkpoint"]:
|
||||
"""Access the latest reported checkpoint to resume from if one exists.
|
||||
|
||||
See :ref:`the checkpoint loading guide <train-dl-loading-checkpoints>` for more details.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import tempfile
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
def train_func(config):
|
||||
start_epoch = 0
|
||||
checkpoint = ray.train.get_checkpoint()
|
||||
if checkpoint:
|
||||
with checkpoint.as_directory() as checkpoint_dir:
|
||||
# Load back training state
|
||||
...
|
||||
|
||||
for epoch in range(start_epoch, config.get("num_epochs", 10)):
|
||||
# Do training...
|
||||
|
||||
metrics = {"loss": ...}
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
# Save the checkpoint...
|
||||
|
||||
checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
|
||||
ray.train.report(metrics, checkpoint=checkpoint)
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_func, scaling_config=ray.train.ScalingConfig(num_workers=2)
|
||||
)
|
||||
|
||||
Returns:
|
||||
Checkpoint object if the session is currently being resumed.
|
||||
Otherwise, return None.
|
||||
"""
|
||||
return get_train_fn_utils().get_checkpoint()
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
@requires_train_worker()
|
||||
def get_all_reported_checkpoints(
|
||||
consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
|
||||
timeout_s: Optional[float] = None,
|
||||
) -> List["ReportedCheckpoint"]:
|
||||
"""Get all the reported checkpoints so far.
|
||||
|
||||
Blocks until Ray Train has finished processing every in-flight `ray.train.report` call.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import tempfile
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
def train_func(config):
|
||||
start_epoch = 0
|
||||
|
||||
for epoch in range(start_epoch, config.get("num_epochs", 2)):
|
||||
# Do training...
|
||||
|
||||
metrics = {"loss": 0.1}
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
# Save the checkpoint...
|
||||
|
||||
checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
|
||||
ray.train.report(metrics, checkpoint=checkpoint)
|
||||
|
||||
reported_checkpoints = ray.train.get_all_reported_checkpoints()
|
||||
# Report artifacts/metrics to experiment tracking framework...
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_func, scaling_config=ray.train.ScalingConfig(num_workers=2)
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Args:
|
||||
consistency_mode: Read semantics for checkpoint retrieval during an ongoing run.
|
||||
Defaults to CheckpointConsistencyMode.VALIDATED.
|
||||
See :class:`~ray.train.CheckpointConsistencyMode` for more details.
|
||||
timeout_s: Timeout in seconds to collecting checkpoint and validation information.
|
||||
Defaults to None to wait indefinitely.
|
||||
|
||||
Returns:
|
||||
List of ReportedCheckpoint objects that represent the checkpoints and
|
||||
corresponding metrics reported by the workers.
|
||||
"""
|
||||
return get_train_fn_utils().get_all_reported_checkpoints(
|
||||
consistency_mode=consistency_mode, timeout_s=timeout_s
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@requires_train_worker()
|
||||
def get_dataset_shard(dataset_name: Optional[str] = None) -> Optional["DataIterator"]:
|
||||
"""Returns the :class:`ray.data.DataIterator` shard for this worker.
|
||||
|
||||
Call :meth:`~ray.data.DataIterator.iter_torch_batches` or
|
||||
:meth:`~ray.data.DataIterator.to_tf` on this shard to convert it to the
|
||||
appropriate framework-specific data type.
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
...
|
||||
for epoch in range(2):
|
||||
# Trainer will automatically handle sharding.
|
||||
data_shard = ray.train.get_dataset_shard("train")
|
||||
for batch in data_shard.iter_torch_batches():
|
||||
...
|
||||
|
||||
train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
datasets={"train": train_dataset}
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Args:
|
||||
dataset_name: If a Dictionary of Datasets was passed to ``Trainer``, then
|
||||
specifies which dataset shard to return.
|
||||
|
||||
Returns:
|
||||
The ``DataIterator`` shard to use for this worker.
|
||||
If no dataset is passed into Trainer, then return None.
|
||||
"""
|
||||
train_fn_utils = get_train_fn_utils()
|
||||
return train_fn_utils.get_dataset_shard(
|
||||
DatasetShardMetadata(
|
||||
dataset_name=dataset_name,
|
||||
world_rank=train_fn_utils.get_context().get_world_rank(),
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,74 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol
|
||||
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train import Checkpoint
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class ValidationFn(Protocol):
|
||||
"""Protocol for a function that validates a checkpoint."""
|
||||
|
||||
def __call__(self, checkpoint: "Checkpoint", **kwargs: Any) -> Dict:
|
||||
...
|
||||
|
||||
|
||||
@dataclass
|
||||
@PublicAPI(stability="alpha")
|
||||
class ValidationTaskConfig:
|
||||
"""Configuration for a specific validation task, passed to report().
|
||||
|
||||
Args:
|
||||
fn_kwargs: json-serializable keyword arguments to pass to the validation function.
|
||||
Note that we always pass `checkpoint` as the first argument to the
|
||||
validation function.
|
||||
timeout_s: Timeout in seconds for this validation task.
|
||||
``None`` is no timeout.
|
||||
"""
|
||||
|
||||
fn_kwargs: Optional[Dict[str, Any]] = None
|
||||
timeout_s: Optional[float] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.fn_kwargs is None:
|
||||
self.fn_kwargs = {}
|
||||
assert (
|
||||
self.timeout_s is None or self.timeout_s > 0
|
||||
), f"The `timeout_s` should be None or greater than zero, actual value: {self.timeout_s}"
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class ValidationConfig:
|
||||
"""Configuration for validation, passed to the trainer.
|
||||
|
||||
Args:
|
||||
fn: The validation function to run on checkpoints.
|
||||
This function should accept a checkpoint as the first argument
|
||||
and return a dictionary of metrics.
|
||||
task_config: Default configuration for validation tasks.
|
||||
The fn_kwargs in this config can be overridden by
|
||||
ValidationTaskConfig passed to report().
|
||||
ray_remote_kwargs: Keyword arguments to pass to `ray.remote()` for the validation task.
|
||||
This can be used to specify resource requirements, number of retries, etc.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fn: ValidationFn,
|
||||
task_config: Optional[ValidationTaskConfig] = None,
|
||||
ray_remote_kwargs: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
self.fn = fn
|
||||
if task_config is None:
|
||||
self.task_config = ValidationTaskConfig()
|
||||
else:
|
||||
self.task_config = task_config
|
||||
# TODO: ray_remote_kwargs is not json-serializable because retry_exceptions
|
||||
# can be a list of exception types. If ray core makes ray_remote_kwargs json-serializable
|
||||
# we can move this to ValidationTaskConfig.
|
||||
if ray_remote_kwargs is None:
|
||||
self.ray_remote_kwargs = {}
|
||||
else:
|
||||
self.ray_remote_kwargs = ray_remote_kwargs
|
||||
Reference in New Issue
Block a user