chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,954 @@
|
||||
import abc
|
||||
import copy
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
|
||||
|
||||
import pyarrow.fs
|
||||
|
||||
import ray
|
||||
import ray.cloudpickle as pickle
|
||||
from ray._common.usage import usage_lib
|
||||
from ray._private.dict import deep_update
|
||||
from ray.air._internal import usage as air_usage
|
||||
from ray.air._internal.config import ensure_only_allowed_dataclass_keys_updated
|
||||
from ray.air._internal.usage import AirEntrypoint
|
||||
from ray.air.config import RunConfig, ScalingConfig
|
||||
from ray.air.result import Result
|
||||
from ray.train import Checkpoint
|
||||
from ray.train._internal.session import get_session
|
||||
from ray.train._internal.storage import (
|
||||
StorageContext,
|
||||
_exists_at_fs_path,
|
||||
get_fs_and_path,
|
||||
)
|
||||
from ray.train.constants import (
|
||||
V2_MIGRATION_GUIDE_MESSAGE,
|
||||
_v2_migration_warnings_enabled,
|
||||
)
|
||||
from ray.train.context import _GET_METADATA_DEPRECATION_MESSAGE
|
||||
from ray.train.utils import _log_deprecation_warning
|
||||
from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import Dataset
|
||||
from ray.tune import Trainable
|
||||
|
||||
_TRAINER_PKL = "trainer.pkl"
|
||||
|
||||
# A type representing either a ray.data.Dataset or a function that returns a
|
||||
# ray.data.Dataset and accepts no arguments.
|
||||
GenDataset = Union["Dataset", Callable[[], "Dataset"]]
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PREPROCESSOR_DEPRECATION_MESSAGE = (
|
||||
"The `preprocessor` argument to Trainers is deprecated as of Ray 2.7. "
|
||||
"Instead, use the Preprocessor `fit` and `transform` APIs directly on the Ray "
|
||||
"Dataset. For any state that needs to be saved to the trained checkpoint, pass it "
|
||||
"in using the `metadata` argument of the `Trainer`. "
|
||||
"For a full example, see "
|
||||
"https://docs.ray.io/en/master/train/user-guides/data-loading-preprocessing.html#preprocessing-structured-data " # noqa:E501
|
||||
)
|
||||
|
||||
_TRAINER_RESTORE_DEPRECATION_WARNING = (
|
||||
"The `restore` and `can_restore` APIs are deprecated and "
|
||||
f"will be removed in a future release. {V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING = (
|
||||
"`resume_from_checkpoint` is deprecated and will be removed in an upcoming "
|
||||
f"release. {V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class TrainingFailedError(RuntimeError):
|
||||
"""An error indicating that training has failed."""
|
||||
|
||||
_RESTORE_MSG = (
|
||||
"The Ray Train run failed. Please inspect the previous error messages for a "
|
||||
"cause. After fixing the issue (assuming that the error is not caused by "
|
||||
"your own application logic, but rather an error such as OOM), you can restart "
|
||||
"the run from scratch or continue this run.\n"
|
||||
"To continue this run, you can use: "
|
||||
'`trainer = {trainer_cls_name}.restore("{path}")`.'
|
||||
)
|
||||
|
||||
_FAILURE_CONFIG_MSG = (
|
||||
"To start a new run that will retry on training failures, set "
|
||||
"`train.RunConfig(failure_config=train.FailureConfig(max_failures))` "
|
||||
"in the Trainer's `run_config` with `max_failures > 0`, or `max_failures = -1` "
|
||||
"for unlimited retries."
|
||||
)
|
||||
|
||||
|
||||
def _train_coordinator_fn(
|
||||
config: dict, trainer_cls: Type["BaseTrainer"], metadata: dict
|
||||
):
|
||||
"""This is the function that defines the logic of the Ray Train coordinator.
|
||||
This is responsible for setting up a remote instance of the `trainer_cls`
|
||||
(a different instance than the one calling `trainer.fit` on the driver!)
|
||||
and running the training loop.
|
||||
"""
|
||||
assert metadata is not None, metadata
|
||||
# Propagate user metadata from the Trainer constructor.
|
||||
get_session().metadata = metadata
|
||||
|
||||
# config already contains merged values.
|
||||
# Instantiate new Trainer in Trainable.
|
||||
trainer = trainer_cls(**config)
|
||||
|
||||
# Get the checkpoint from Tune and pass it to workers later on.
|
||||
checkpoint = ray.tune.get_checkpoint()
|
||||
if checkpoint:
|
||||
# Set `starting_checkpoint` for auto-recovery fault-tolerance
|
||||
# as well as manual restoration.
|
||||
trainer.starting_checkpoint = checkpoint
|
||||
# else: Train will restore from the user-provided
|
||||
# `resume_from_checkpoint` == `starting_checkpoint`.
|
||||
|
||||
# Evaluate datasets if they are wrapped in a factory.
|
||||
trainer.datasets = {
|
||||
k: d() if callable(d) else d for k, d in trainer.datasets.items()
|
||||
}
|
||||
|
||||
trainer.setup()
|
||||
trainer.training_loop()
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class BaseTrainer(abc.ABC):
|
||||
"""Defines interface for distributed training on Ray.
|
||||
|
||||
Note: The base ``BaseTrainer`` class cannot be instantiated directly. Only
|
||||
one of its subclasses can be used.
|
||||
|
||||
Note to developers: If a new trainer is added, please update
|
||||
`air/_internal/usage.py`.
|
||||
|
||||
**How does a trainer work?**
|
||||
|
||||
- First, initialize the Trainer. The initialization runs locally,
|
||||
so heavyweight setup should not be done in ``__init__``.
|
||||
- Then, when you call ``trainer.fit()``, the Trainer is serialized
|
||||
and copied to a remote Ray actor. The following methods are then
|
||||
called in sequence on the remote actor.
|
||||
- ``trainer.setup()``: Any heavyweight Trainer setup should be
|
||||
specified here.
|
||||
- ``trainer.training_loop()``: Executes the main training logic.
|
||||
- Calling ``trainer.fit()`` will return a ``ray.result.Result``
|
||||
object where you can access metrics from your training run, as well
|
||||
as any checkpoints that may have been saved.
|
||||
|
||||
**How do I create a new Trainer?**
|
||||
|
||||
Subclass ``ray.train.trainer.BaseTrainer``, and override the ``training_loop``
|
||||
method, and optionally ``setup``.
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import torch
|
||||
|
||||
from ray.train.trainer import BaseTrainer
|
||||
from ray import train, tune
|
||||
|
||||
|
||||
class MyPytorchTrainer(BaseTrainer):
|
||||
def setup(self):
|
||||
self.model = torch.nn.Linear(1, 1)
|
||||
self.optimizer = torch.optim.SGD(
|
||||
self.model.parameters(), lr=0.1)
|
||||
|
||||
def training_loop(self):
|
||||
# You can access any Trainer attributes directly in this method.
|
||||
# self.datasets["train"] has already been
|
||||
dataset = self.datasets["train"]
|
||||
|
||||
torch_ds = dataset.iter_torch_batches(dtypes=torch.float)
|
||||
loss_fn = torch.nn.MSELoss()
|
||||
|
||||
for epoch_idx in range(10):
|
||||
loss = 0
|
||||
num_batches = 0
|
||||
torch_ds = dataset.iter_torch_batches(
|
||||
dtypes=torch.float, batch_size=2
|
||||
)
|
||||
for batch in torch_ds:
|
||||
X = torch.unsqueeze(batch["x"], 1)
|
||||
y = torch.unsqueeze(batch["y"], 1)
|
||||
# Compute prediction error
|
||||
pred = self.model(X)
|
||||
batch_loss = loss_fn(pred, y)
|
||||
|
||||
# Backpropagation
|
||||
self.optimizer.zero_grad()
|
||||
batch_loss.backward()
|
||||
self.optimizer.step()
|
||||
|
||||
loss += batch_loss.item()
|
||||
num_batches += 1
|
||||
loss /= num_batches
|
||||
|
||||
# Use Tune functions to report intermediate
|
||||
# results.
|
||||
train.report({"loss": loss, "epoch": epoch_idx})
|
||||
|
||||
|
||||
# Initialize the Trainer, and call Trainer.fit()
|
||||
import ray
|
||||
train_dataset = ray.data.from_items(
|
||||
[{"x": i, "y": i} for i in range(10)])
|
||||
my_trainer = MyPytorchTrainer(datasets={"train": train_dataset})
|
||||
result = my_trainer.fit()
|
||||
|
||||
Args:
|
||||
scaling_config: Configuration for how to scale training.
|
||||
run_config: Configuration for the execution of the training run.
|
||||
datasets: Any Datasets to use for training. Use the key "train"
|
||||
to denote which dataset is the training dataset.
|
||||
metadata: Dict that should be made available via
|
||||
`train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
|
||||
for checkpoints saved from this Trainer. Must be JSON-serializable.
|
||||
resume_from_checkpoint: A checkpoint to resume training from.
|
||||
"""
|
||||
|
||||
_scaling_config_allowed_keys: List[str] = [
|
||||
"trainer_resources",
|
||||
]
|
||||
_handles_checkpoint_freq: bool = False
|
||||
_handles_checkpoint_at_end: bool = False
|
||||
|
||||
# fields to propagate to Tuner param_space.
|
||||
# See `BaseTrainer._extract_fields_for_tuner_param_space` for more details.
|
||||
_fields_for_tuner_param_space = []
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
):
|
||||
self.scaling_config = (
|
||||
scaling_config if scaling_config is not None else ScalingConfig()
|
||||
)
|
||||
self.run_config = (
|
||||
copy.copy(run_config) if run_config is not None else RunConfig()
|
||||
)
|
||||
self.metadata = metadata
|
||||
self.datasets = datasets if datasets is not None else {}
|
||||
self.starting_checkpoint = resume_from_checkpoint
|
||||
|
||||
if _v2_migration_warnings_enabled():
|
||||
if metadata is not None:
|
||||
_log_deprecation_warning(_GET_METADATA_DEPRECATION_MESSAGE)
|
||||
if resume_from_checkpoint is not None:
|
||||
_log_deprecation_warning(_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING)
|
||||
|
||||
# These attributes should only be set through `BaseTrainer.restore`
|
||||
self._restore_path = None
|
||||
self._restore_storage_filesystem = None
|
||||
|
||||
self._validate_attributes()
|
||||
|
||||
usage_lib.record_library_usage("train")
|
||||
air_usage.tag_air_trainer(self)
|
||||
|
||||
@classmethod
|
||||
@Deprecated(message=_TRAINER_RESTORE_DEPRECATION_WARNING)
|
||||
def restore(
|
||||
cls: Type["BaseTrainer"],
|
||||
path: Union[str, os.PathLike],
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
**kwargs,
|
||||
) -> "BaseTrainer":
|
||||
"""Restores a Train experiment from a previously interrupted/failed run.
|
||||
|
||||
Restore should be used for experiment-level fault tolerance in the event
|
||||
that the head node crashes (e.g., OOM or some other runtime error) or the
|
||||
entire cluster goes down (e.g., network error affecting all nodes).
|
||||
|
||||
A run that has already completed successfully will not be resumed from this API.
|
||||
To continue training from a successful run, launch a new run with the
|
||||
``<Framework>Trainer(resume_from_checkpoint)`` API instead, passing in a
|
||||
checkpoint from the previous run to start with.
|
||||
|
||||
.. warning::
|
||||
|
||||
The ``path`` must point to a **trusted** experiment directory.
|
||||
Restoring from an untrusted path executes arbitrary Python code
|
||||
(the experiment state uses pickle serialization). Never restore
|
||||
from a path that other parties can write to.
|
||||
|
||||
.. note::
|
||||
|
||||
Restoring an experiment from a path that's pointing to a *different*
|
||||
location than the original experiment path is supported. However, Ray Train
|
||||
assumes that the full experiment directory is available
|
||||
(including checkpoints) so that it's possible to resume trials from their
|
||||
latest state.
|
||||
|
||||
For example, if the original experiment path was run locally, then the
|
||||
results are uploaded to cloud storage, Ray Train expects the full contents
|
||||
to be available in cloud storage if attempting to resume
|
||||
via ``<Framework>Trainer.restore("s3://...")``. The restored run will
|
||||
continue writing results to the same cloud storage location.
|
||||
|
||||
The following example can be paired with implementing job retry using
|
||||
:ref:`Ray Jobs <jobs-overview>` to produce a Train experiment that will
|
||||
attempt to resume on both experiment-level and trial-level failures:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import os
|
||||
import ray
|
||||
from ray import train
|
||||
from ray.train.trainer import BaseTrainer
|
||||
|
||||
experiment_name = "unique_experiment_name"
|
||||
storage_path = os.path.expanduser("~/ray_results")
|
||||
experiment_dir = os.path.join(storage_path, experiment_name)
|
||||
|
||||
# Define some dummy inputs for demonstration purposes
|
||||
datasets = {"train": ray.data.from_items([{"a": i} for i in range(10)])}
|
||||
|
||||
class CustomTrainer(BaseTrainer):
|
||||
def training_loop(self):
|
||||
pass
|
||||
|
||||
if CustomTrainer.can_restore(experiment_dir):
|
||||
trainer = CustomTrainer.restore(
|
||||
experiment_dir, datasets=datasets
|
||||
)
|
||||
else:
|
||||
trainer = CustomTrainer(
|
||||
datasets=datasets,
|
||||
run_config=train.RunConfig(
|
||||
name=experiment_name,
|
||||
storage_path=storage_path,
|
||||
# Tip: You can also enable retries on failure for
|
||||
# worker-level fault tolerance
|
||||
failure_config=train.FailureConfig(max_failures=3),
|
||||
),
|
||||
)
|
||||
|
||||
result = trainer.fit()
|
||||
|
||||
Args:
|
||||
path: The path to the experiment directory of the training run to restore.
|
||||
This can be a local path or a remote URI if the experiment was
|
||||
uploaded to the cloud.
|
||||
storage_filesystem: Custom ``pyarrow.fs.FileSystem``
|
||||
corresponding to the ``path``. This may be necessary if the original
|
||||
experiment passed in a custom filesystem.
|
||||
datasets: Re-specified datasets used in the original training run.
|
||||
This must include all the datasets that were passed in the
|
||||
original trainer constructor.
|
||||
scaling_config: Optionally re-specified scaling config. This can be
|
||||
modified to be different from the original spec.
|
||||
**kwargs: Other optionally re-specified arguments, passed in by subclasses.
|
||||
|
||||
Raises:
|
||||
ValueError: If all datasets were not re-supplied on restore.
|
||||
|
||||
Returns:
|
||||
BaseTrainer: A restored instance of the class that is calling this method.
|
||||
"""
|
||||
if _v2_migration_warnings_enabled():
|
||||
_log_deprecation_warning(_TRAINER_RESTORE_DEPRECATION_WARNING)
|
||||
|
||||
if not cls.can_restore(path, storage_filesystem):
|
||||
raise ValueError(
|
||||
f"Invalid restore path: {path}. Make sure that this path exists and "
|
||||
"is the experiment directory that results from a call to "
|
||||
"`trainer.fit()`."
|
||||
)
|
||||
fs, fs_path = get_fs_and_path(path, storage_filesystem)
|
||||
trainer_pkl_path = Path(fs_path, _TRAINER_PKL).as_posix()
|
||||
with fs.open_input_file(trainer_pkl_path) as f:
|
||||
trainer_cls, param_dict = pickle.loads(f.readall())
|
||||
|
||||
if trainer_cls is not cls:
|
||||
warnings.warn(
|
||||
f"Invalid trainer type. You are attempting to restore a trainer of type"
|
||||
f" {trainer_cls} with `{cls.__name__}.restore`, "
|
||||
"which will most likely fail. "
|
||||
f"Use `{trainer_cls.__name__}.restore` instead."
|
||||
)
|
||||
|
||||
original_datasets = param_dict.pop("datasets", {})
|
||||
if original_datasets and not datasets:
|
||||
raise ValueError(
|
||||
"The following datasets need to be provided again on restore: "
|
||||
f"{list(original_datasets.keys())}\n"
|
||||
f"Use {cls.__name__}.restore(..., datasets=datasets) "
|
||||
"with the datasets that were provided to the original trainer."
|
||||
)
|
||||
datasets = datasets or {}
|
||||
if set(original_datasets) != set(datasets):
|
||||
raise ValueError(
|
||||
"The provided datasets don't match the original dataset keys.\n"
|
||||
f" Expected datasets for the keys: {list(original_datasets.keys())}\n"
|
||||
f" Actual datasets provided: {list(datasets.keys())}"
|
||||
)
|
||||
param_dict["datasets"] = datasets
|
||||
|
||||
if scaling_config:
|
||||
param_dict["scaling_config"] = scaling_config
|
||||
|
||||
for param_name, val in kwargs.items():
|
||||
# Overwrite the old value if something is passed into restore
|
||||
if val is not None:
|
||||
param_dict[param_name] = val
|
||||
|
||||
try:
|
||||
trainer = cls(**param_dict)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
"Trainer restoration failed (see above for the stack trace). "
|
||||
"Make sure that you use the right trainer class to restore: "
|
||||
f"`{cls.__name__}.restore`\n"
|
||||
) from e
|
||||
trainer._restore_path = path
|
||||
trainer._restore_storage_filesystem = storage_filesystem
|
||||
return trainer
|
||||
|
||||
@classmethod
|
||||
@Deprecated(
|
||||
message=_TRAINER_RESTORE_DEPRECATION_WARNING,
|
||||
warning=_v2_migration_warnings_enabled(),
|
||||
)
|
||||
def can_restore(
|
||||
cls: Type["BaseTrainer"],
|
||||
path: Union[str, os.PathLike],
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
||||
) -> bool:
|
||||
"""Checks whether a given directory contains a restorable Train experiment.
|
||||
|
||||
Args:
|
||||
path: The path to the experiment directory of the Train experiment.
|
||||
This can be either a local directory (e.g., ~/ray_results/exp_name)
|
||||
or a remote URI (e.g., s3://bucket/exp_name).
|
||||
storage_filesystem: Custom ``pyarrow.fs.FileSystem`` to use. If not
|
||||
provided, the filesystem is auto-resolved from ``path``.
|
||||
|
||||
Returns:
|
||||
bool: Whether this path exists and contains the trainer state to resume from
|
||||
"""
|
||||
if _v2_migration_warnings_enabled():
|
||||
_log_deprecation_warning(_TRAINER_RESTORE_DEPRECATION_WARNING)
|
||||
|
||||
fs, fs_path = get_fs_and_path(path, storage_filesystem)
|
||||
trainer_pkl_path = Path(fs_path, _TRAINER_PKL).as_posix()
|
||||
return _exists_at_fs_path(fs, trainer_pkl_path)
|
||||
|
||||
def __repr__(self):
|
||||
# A dictionary that maps parameters to their default values.
|
||||
default_values: Dict[str, Any] = {
|
||||
"scaling_config": ScalingConfig(),
|
||||
"run_config": RunConfig(),
|
||||
"datasets": {},
|
||||
"starting_checkpoint": None,
|
||||
}
|
||||
|
||||
non_default_arguments = []
|
||||
for parameter, default_value in default_values.items():
|
||||
value = getattr(self, parameter)
|
||||
if value != default_value:
|
||||
# 'Dataset.__repr__' returns a table rather than a regular Python object
|
||||
# representation. So, we need to special case the 'datasets' parameter.
|
||||
if parameter == "datasets":
|
||||
value_repr = format_datasets_for_repr(value)
|
||||
else:
|
||||
value_repr = repr(value)
|
||||
|
||||
non_default_arguments.append(f"{parameter}={value_repr}")
|
||||
|
||||
if non_default_arguments:
|
||||
return f"<{self.__class__.__name__} {' '.join(non_default_arguments)}>"
|
||||
|
||||
return f"<{self.__class__.__name__}>"
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
# Store the init args as attributes so this can be merged with Tune hparams.
|
||||
trainer = super(BaseTrainer, cls).__new__(cls)
|
||||
parameters = inspect.signature(cls.__init__).parameters
|
||||
parameters = list(parameters.keys())
|
||||
# Remove self.
|
||||
parameters = parameters[1:]
|
||||
arg_dict = dict(zip(parameters, args))
|
||||
trainer._param_dict = {**arg_dict, **kwargs}
|
||||
return trainer
|
||||
|
||||
def _validate_attributes(self):
|
||||
"""Called on __init()__ to validate trainer attributes."""
|
||||
# Run config
|
||||
if not isinstance(self.run_config, RunConfig):
|
||||
raise ValueError(
|
||||
f"`run_config` should be an instance of `ray.train.RunConfig`, "
|
||||
f"found {type(self.run_config)} with value `{self.run_config}`."
|
||||
)
|
||||
# Scaling config
|
||||
if not isinstance(self.scaling_config, ScalingConfig):
|
||||
raise ValueError(
|
||||
"`scaling_config` should be an instance of `ScalingConfig`, "
|
||||
f"found {type(self.scaling_config)} with value `{self.scaling_config}`."
|
||||
)
|
||||
# Datasets
|
||||
if not isinstance(self.datasets, dict):
|
||||
raise ValueError(
|
||||
f"`datasets` should be a dict mapping from a string to "
|
||||
f"`ray.data.Dataset` objects, "
|
||||
f"found {type(self.datasets)} with value `{self.datasets}`."
|
||||
)
|
||||
else:
|
||||
for key, dataset in self.datasets.items():
|
||||
if not isinstance(dataset, ray.data.Dataset) and not callable(dataset):
|
||||
raise ValueError(
|
||||
f"The Dataset under '{key}' key is not a "
|
||||
"`ray.data.Dataset`. "
|
||||
f"Received {dataset} instead."
|
||||
)
|
||||
# Metadata.
|
||||
self.metadata = self.metadata or {}
|
||||
if not isinstance(self.metadata, dict):
|
||||
raise TypeError(
|
||||
f"The provided metadata must be a dict, was {type(self.metadata)}."
|
||||
)
|
||||
try:
|
||||
self.metadata = json.loads(json.dumps(self.metadata))
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
"The provided metadata must be JSON-serializable: "
|
||||
f"{self.metadata}: {e}"
|
||||
)
|
||||
|
||||
if self.starting_checkpoint is not None and not isinstance(
|
||||
self.starting_checkpoint, Checkpoint
|
||||
):
|
||||
raise ValueError(
|
||||
f"`resume_from_checkpoint` should be an instance of "
|
||||
f"`ray.train.Checkpoint`, found {type(self.starting_checkpoint)} "
|
||||
f"with value `{self.starting_checkpoint}`."
|
||||
)
|
||||
|
||||
self._log_v2_deprecation_warnings()
|
||||
|
||||
def _log_v2_deprecation_warnings(self):
|
||||
"""Logs deprecation warnings for v2 migration.
|
||||
|
||||
Log them here in the Ray Train case rather than in the configuration
|
||||
constructors to avoid logging incorrect deprecation warnings when
|
||||
`ray.train.RunConfig` is passed to Ray Tune.
|
||||
"""
|
||||
from ray.train.v2._internal.constants import V2_ENABLED_ENV_VAR, is_v2_enabled
|
||||
|
||||
if is_v2_enabled():
|
||||
raise DeprecationWarning(
|
||||
f"Detected use of a deprecated Trainer import from `{self.__class__.__module__}`. "
|
||||
"This Trainer class is not compatible with Ray Train V2.\n"
|
||||
"To fix this:\n"
|
||||
" - Update to use the new import path. For example, "
|
||||
"`from ray.train.torch.torch_trainer import TorchTrainer` -> "
|
||||
"`from ray.train.torch import TorchTrainer`\n"
|
||||
f" - Or, explicitly disable V2 by setting: {V2_ENABLED_ENV_VAR}=0\n"
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
if not _v2_migration_warnings_enabled():
|
||||
return
|
||||
|
||||
from ray.train.v2._internal.migration_utils import (
|
||||
CALLBACKS_DEPRECATION_MESSAGE,
|
||||
FAIL_FAST_DEPRECATION_MESSAGE,
|
||||
LOG_TO_FILE_DEPRECATION_MESSAGE,
|
||||
PROGRESS_REPORTER_DEPRECATION_MESSAGE,
|
||||
STOP_DEPRECATION_MESSAGE,
|
||||
SYNC_CONFIG_DEPRECATION_MESSAGE,
|
||||
TRAINER_RESOURCES_DEPRECATION_MESSAGE,
|
||||
VERBOSE_DEPRECATION_MESSAGE,
|
||||
)
|
||||
|
||||
# ScalingConfig deprecations
|
||||
if self.scaling_config.trainer_resources is not None:
|
||||
_log_deprecation_warning(TRAINER_RESOURCES_DEPRECATION_MESSAGE)
|
||||
|
||||
# FailureConfig deprecations
|
||||
if self.run_config.failure_config.fail_fast:
|
||||
_log_deprecation_warning(FAIL_FAST_DEPRECATION_MESSAGE)
|
||||
|
||||
# RunConfig deprecations
|
||||
# NOTE: _verbose is the original verbose value passed by the user
|
||||
if self.run_config._verbose is not None:
|
||||
_log_deprecation_warning(VERBOSE_DEPRECATION_MESSAGE)
|
||||
|
||||
if self.run_config.log_to_file:
|
||||
_log_deprecation_warning(LOG_TO_FILE_DEPRECATION_MESSAGE)
|
||||
|
||||
if self.run_config.stop is not None:
|
||||
_log_deprecation_warning(STOP_DEPRECATION_MESSAGE)
|
||||
|
||||
if self.run_config.callbacks is not None:
|
||||
_log_deprecation_warning(CALLBACKS_DEPRECATION_MESSAGE)
|
||||
|
||||
if self.run_config.progress_reporter is not None:
|
||||
_log_deprecation_warning(PROGRESS_REPORTER_DEPRECATION_MESSAGE)
|
||||
|
||||
if self.run_config.sync_config != ray.train.SyncConfig():
|
||||
_log_deprecation_warning(SYNC_CONFIG_DEPRECATION_MESSAGE)
|
||||
|
||||
@classmethod
|
||||
def _validate_scaling_config(cls, scaling_config: ScalingConfig) -> ScalingConfig:
|
||||
"""Returns scaling config dataclass after validating updated keys."""
|
||||
ensure_only_allowed_dataclass_keys_updated(
|
||||
dataclass=scaling_config,
|
||||
allowed_keys=cls._scaling_config_allowed_keys,
|
||||
)
|
||||
return scaling_config
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Called during fit() to perform initial setup on the Trainer.
|
||||
|
||||
.. note:: This method is run on a remote process.
|
||||
|
||||
This method will not be called on the driver, so any expensive setup
|
||||
operations should be placed here and not in ``__init__``.
|
||||
|
||||
This method is called prior to ``preprocess_datasets`` and
|
||||
``training_loop``.
|
||||
"""
|
||||
pass
|
||||
|
||||
def preprocess_datasets(self) -> None:
|
||||
"""Deprecated."""
|
||||
raise DeprecationWarning(
|
||||
"`preprocess_datasets` is no longer used, since preprocessors "
|
||||
f"are no longer accepted by Trainers.\n{PREPROCESSOR_DEPRECATION_MESSAGE}"
|
||||
)
|
||||
|
||||
@abc.abstractmethod
|
||||
def training_loop(self) -> None:
|
||||
"""Loop called by fit() to run training and report results to Tune.
|
||||
|
||||
.. note:: This method runs on a remote process.
|
||||
|
||||
``self.datasets`` have already been evaluated if they were wrapped in a factory.
|
||||
|
||||
You can use the :ref:`Ray Train utilities <train-loop-api>`
|
||||
(:func:`train.report() <ray.train.report>` and
|
||||
:func:`train.get_checkpoint() <ray.train.get_checkpoint>`) inside
|
||||
this training loop.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
from ray.train.trainer import BaseTrainer
|
||||
from ray import train
|
||||
|
||||
class MyTrainer(BaseTrainer):
|
||||
def training_loop(self):
|
||||
for epoch_idx in range(5):
|
||||
...
|
||||
train.report({"epoch": epoch_idx})
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
def fit(self) -> Result:
|
||||
"""Runs training.
|
||||
|
||||
Returns:
|
||||
A Result object containing the training result.
|
||||
|
||||
Raises:
|
||||
ray.train.base_trainer.TrainingFailedError: If any failures during the execution
|
||||
of ``self.as_trainable()``, or during the Tune execution loop.
|
||||
"""
|
||||
from ray.tune import ResumeConfig, TuneError
|
||||
from ray.tune.tuner import Tuner
|
||||
|
||||
trainable = self.as_trainable()
|
||||
param_space = self._extract_fields_for_tuner_param_space()
|
||||
|
||||
self.run_config.name = (
|
||||
self.run_config.name or StorageContext.get_experiment_dir_name(trainable)
|
||||
)
|
||||
# The storage context here is only used to access the resolved
|
||||
# storage fs and experiment path, in order to avoid duplicating that logic.
|
||||
# This is NOT the storage context object that gets passed to remote workers.
|
||||
storage = StorageContext(
|
||||
storage_path=self.run_config.storage_path,
|
||||
experiment_dir_name=self.run_config.name,
|
||||
storage_filesystem=self.run_config.storage_filesystem,
|
||||
)
|
||||
|
||||
if self._restore_path:
|
||||
tuner = Tuner.restore(
|
||||
path=self._restore_path,
|
||||
trainable=trainable,
|
||||
param_space=param_space,
|
||||
_resume_config=ResumeConfig(
|
||||
finished=ResumeConfig.ResumeType.RESUME,
|
||||
unfinished=ResumeConfig.ResumeType.RESUME,
|
||||
errored=ResumeConfig.ResumeType.RESUME,
|
||||
),
|
||||
storage_filesystem=self._restore_storage_filesystem,
|
||||
)
|
||||
else:
|
||||
tuner = Tuner(
|
||||
trainable=trainable,
|
||||
param_space=param_space,
|
||||
run_config=self.run_config,
|
||||
_entrypoint=AirEntrypoint.TRAINER,
|
||||
)
|
||||
|
||||
self._save(storage.storage_filesystem, storage.experiment_fs_path)
|
||||
|
||||
restore_msg = TrainingFailedError._RESTORE_MSG.format(
|
||||
trainer_cls_name=self.__class__.__name__,
|
||||
path=str(storage.experiment_fs_path),
|
||||
)
|
||||
|
||||
try:
|
||||
result_grid = tuner.fit()
|
||||
except TuneError as e:
|
||||
# Catch any `TuneError`s raised by the `Tuner.fit` call.
|
||||
# Unwrap the `TuneError` if needed.
|
||||
parent_error = e.__cause__ or e
|
||||
|
||||
# Raise it to the user as a `TrainingFailedError` with a message to restore.
|
||||
raise TrainingFailedError(restore_msg) from parent_error
|
||||
# Other exceptions get passed through directly (ex: on `fail_fast='raise'`)
|
||||
|
||||
assert len(result_grid) == 1
|
||||
result = result_grid[0]
|
||||
if result.error:
|
||||
# Raise trainable errors to the user with a message to restore
|
||||
# or configure `FailureConfig` in a new run.
|
||||
raise TrainingFailedError(
|
||||
"\n".join([restore_msg, TrainingFailedError._FAILURE_CONFIG_MSG])
|
||||
) from result.error
|
||||
return result
|
||||
|
||||
def _save(self, fs: pyarrow.fs.FileSystem, experiment_path: str):
|
||||
"""Saves the current trainer's class along with the `param_dict` of
|
||||
parameters passed to this trainer's constructor.
|
||||
|
||||
This is used to recreate the trainer on restore.
|
||||
Unless a parameter is re-specified during restoration (only a subset
|
||||
of parameters can be passed in again), that parameter will be loaded
|
||||
from the saved copy.
|
||||
|
||||
Datasets should not be saved as part of the state. Instead, we save the
|
||||
keys and replace the dataset values with dummy functions that will
|
||||
raise an error if invoked. The error only serves as a guardrail for
|
||||
misuse (e.g., manually unpickling and constructing the Trainer again)
|
||||
and is not typically surfaced, since datasets must be re-specified
|
||||
upon restoration.
|
||||
"""
|
||||
param_dict = self._param_dict.copy()
|
||||
datasets = param_dict.pop("datasets", {})
|
||||
|
||||
def raise_fn():
|
||||
raise RuntimeError
|
||||
|
||||
if datasets:
|
||||
param_dict["datasets"] = dict.fromkeys(datasets, raise_fn)
|
||||
|
||||
cls_and_param_dict = (self.__class__, param_dict)
|
||||
|
||||
fs.create_dir(experiment_path)
|
||||
with fs.open_output_stream(Path(experiment_path, _TRAINER_PKL).as_posix()) as f:
|
||||
f.write(pickle.dumps(cls_and_param_dict))
|
||||
|
||||
def _extract_fields_for_tuner_param_space(self) -> Dict:
|
||||
"""Extracts fields to be included in `Tuner.param_space`.
|
||||
|
||||
This is needed to leverage the full logging/integration offerings from Tune.
|
||||
For example, `param_space` is logged automatically to wandb integration.
|
||||
|
||||
Currently only done for `train_loop_config`.
|
||||
|
||||
Returns:
|
||||
A dictionary that should be passed to Tuner.param_space.
|
||||
"""
|
||||
result = {}
|
||||
for key in self._fields_for_tuner_param_space:
|
||||
if key in self._param_dict.keys():
|
||||
result[key] = copy.deepcopy(self._param_dict[key])
|
||||
return result
|
||||
|
||||
def _generate_trainable_cls(self) -> Type["Trainable"]:
|
||||
"""Generates the base Trainable class.
|
||||
|
||||
Returns:
|
||||
A Trainable class to use for training.
|
||||
"""
|
||||
|
||||
from ray.tune.execution.placement_groups import PlacementGroupFactory
|
||||
from ray.tune.trainable import wrap_function
|
||||
|
||||
trainer_cls = self.__class__
|
||||
scaling_config = self.scaling_config
|
||||
metadata = self.metadata
|
||||
|
||||
train_coordinator_fn = partial(
|
||||
_train_coordinator_fn, trainer_cls=trainer_cls, metadata=metadata
|
||||
)
|
||||
# Change the name of the training function to match the name of the Trainer
|
||||
# class. This will mean the Tune trial name will match the name of Trainer on
|
||||
# stdout messages and the results directory.
|
||||
train_coordinator_fn.__name__ = trainer_cls.__name__
|
||||
|
||||
trainable_cls = wrap_function(train_coordinator_fn)
|
||||
has_base_dataset = bool(self.datasets)
|
||||
if has_base_dataset:
|
||||
from ray.data.context import DataContext
|
||||
|
||||
dataset_context = DataContext.get_current()
|
||||
else:
|
||||
dataset_context = None
|
||||
|
||||
class TrainTrainable(trainable_cls):
|
||||
"""Adds default resources to the Trainable."""
|
||||
|
||||
_handles_checkpoint_freq = trainer_cls._handles_checkpoint_freq
|
||||
_handles_checkpoint_at_end = trainer_cls._handles_checkpoint_at_end
|
||||
|
||||
@classmethod
|
||||
def has_base_dataset(cls) -> bool:
|
||||
"""Whether a dataset is provided through the Trainer."""
|
||||
return has_base_dataset
|
||||
|
||||
@classmethod
|
||||
def base_scaling_config(cls) -> ScalingConfig:
|
||||
"""Returns the unchanged scaling config provided through the Trainer."""
|
||||
return scaling_config
|
||||
|
||||
def setup(self, config, **kwargs):
|
||||
base_config = dict(kwargs)
|
||||
# Merge Tuner param space hyperparameters in `config` into the
|
||||
# base config passed to the Trainer constructor, which is `base_config`.
|
||||
# `base_config` is pulled from the object store from the usage of
|
||||
# tune.with_parameters in `BaseTrainer.as_trainable`.
|
||||
|
||||
# run_config is not a tunable hyperparameter so it does not need to be
|
||||
# merged.
|
||||
run_config = base_config.pop("run_config", None)
|
||||
self._merged_config = deep_update(
|
||||
base_config, self.config, new_keys_allowed=True
|
||||
)
|
||||
self._merged_config["run_config"] = run_config
|
||||
merged_scaling_config = self._merged_config.get(
|
||||
"scaling_config", ScalingConfig()
|
||||
)
|
||||
if isinstance(merged_scaling_config, dict):
|
||||
merged_scaling_config = ScalingConfig(**merged_scaling_config)
|
||||
self._merged_config[
|
||||
"scaling_config"
|
||||
] = self._reconcile_scaling_config_with_trial_resources(
|
||||
merged_scaling_config
|
||||
)
|
||||
if self.has_base_dataset():
|
||||
# Set the DataContext on the Trainer actor to the DataContext
|
||||
# specified on the driver.
|
||||
DataContext._set_current(dataset_context)
|
||||
super(TrainTrainable, self).setup(config)
|
||||
|
||||
def _reconcile_scaling_config_with_trial_resources(
|
||||
self, scaling_config: ScalingConfig
|
||||
) -> ScalingConfig:
|
||||
"""
|
||||
ResourceChangingScheduler workaround.
|
||||
|
||||
Ensures that the scaling config matches trial resources.
|
||||
|
||||
This should be replaced with RCS returning a ScalingConfig
|
||||
in the future.
|
||||
"""
|
||||
|
||||
trial_resources = self.trial_resources
|
||||
# This will be false if the resources are default
|
||||
if not isinstance(trial_resources, PlacementGroupFactory):
|
||||
return scaling_config
|
||||
|
||||
# Ignore ResourceChangingScheduler workaround when resource bundles
|
||||
# are unchanged
|
||||
if self.trial_resources == scaling_config.as_placement_group_factory():
|
||||
return scaling_config
|
||||
|
||||
trainer_cls._validate_scaling_config(scaling_config)
|
||||
|
||||
return ScalingConfig.from_placement_group_factory(trial_resources)
|
||||
|
||||
def _trainable_func(self, config):
|
||||
# We ignore the config passed by Tune and instead use the merged
|
||||
# config which includes the initial Trainer args.
|
||||
super()._trainable_func(self._merged_config)
|
||||
|
||||
@classmethod
|
||||
def default_resource_request(cls, config):
|
||||
# `config["scaling_config"] is a dataclass when passed via the
|
||||
# `scaling_config` argument in `Trainer` and is a dict when passed
|
||||
# via the `scaling_config` key of `param_spec`.
|
||||
|
||||
# Conversion logic must be duplicated in `TrainTrainable.__init__`
|
||||
# because this is a class method.
|
||||
updated_scaling_config = config.get("scaling_config", scaling_config)
|
||||
if isinstance(updated_scaling_config, dict):
|
||||
updated_scaling_config = ScalingConfig(**updated_scaling_config)
|
||||
validated_scaling_config = trainer_cls._validate_scaling_config(
|
||||
updated_scaling_config
|
||||
)
|
||||
return validated_scaling_config.as_placement_group_factory()
|
||||
|
||||
return TrainTrainable
|
||||
|
||||
def as_trainable(self) -> Type["Trainable"]:
|
||||
"""Converts self to a ``tune.Trainable`` class."""
|
||||
from ray import tune
|
||||
|
||||
base_config = self._param_dict
|
||||
trainable_cls = self._generate_trainable_cls()
|
||||
|
||||
# Wrap with `tune.with_parameters` to handle very large values in base_config
|
||||
return tune.with_parameters(trainable_cls, **base_config)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
def format_datasets_for_repr(datasets: Optional[Dict[str, GenDataset]]) -> str:
|
||||
"""Format datasets for BaseTrainer repr using plan strings.
|
||||
|
||||
The Dataset.__repr__ returns a table rather than a conventional Python object
|
||||
reprentation. To ensure the BaseTrainer representation still looks reasonable, we
|
||||
need to special-case datasets.
|
||||
"""
|
||||
from ray.data import Dataset
|
||||
from ray.data._internal.dataset_repr import build_dataset_summary_repr
|
||||
|
||||
assert datasets is not None, "Expected caller to pass in non-None argument"
|
||||
|
||||
formatted = {}
|
||||
for key, dataset in datasets.items():
|
||||
if isinstance(dataset, Dataset):
|
||||
formatted[key] = build_dataset_summary_repr(dataset)
|
||||
else:
|
||||
formatted[key] = dataset
|
||||
|
||||
return "{" + ", ".join(f"'{key}': {formatted[key]}" for key in datasets) + "}"
|
||||
Reference in New Issue
Block a user