chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,515 @@
|
||||
import glob
|
||||
import warnings
|
||||
from abc import ABCMeta
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
||||
|
||||
import ray.tune
|
||||
from ray.tune.utils.util import _atomic_save, _load_newest_checkpoint
|
||||
from ray.util.annotations import DeveloperAPI, PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.tune.experiment import Trial
|
||||
from ray.tune.stopper import Stopper
|
||||
|
||||
|
||||
class _CallbackMeta(ABCMeta):
|
||||
"""A helper metaclass to ensure container classes (e.g. CallbackList) have
|
||||
implemented all the callback methods (e.g. `on_*`).
|
||||
"""
|
||||
|
||||
def __new__(mcs, name: str, bases: Tuple[type], attrs: Dict[str, Any]) -> type:
|
||||
cls = super().__new__(mcs, name, bases, attrs)
|
||||
|
||||
if mcs.need_check(cls, name, bases, attrs):
|
||||
mcs.check(cls, name, bases, attrs)
|
||||
|
||||
return cls
|
||||
|
||||
@classmethod
|
||||
def need_check(
|
||||
mcs, cls: type, name: str, bases: Tuple[type], attrs: Dict[str, Any]
|
||||
) -> bool:
|
||||
return attrs.get("IS_CALLBACK_CONTAINER", False)
|
||||
|
||||
@classmethod
|
||||
def check(
|
||||
mcs, cls: type, name: str, bases: Tuple[type], attrs: Dict[str, Any]
|
||||
) -> None:
|
||||
methods = set()
|
||||
for base in bases:
|
||||
methods.update(
|
||||
attr_name
|
||||
for attr_name, attr in vars(base).items()
|
||||
if mcs.need_override_by_subclass(attr_name, attr)
|
||||
)
|
||||
overridden = {
|
||||
attr_name
|
||||
for attr_name, attr in attrs.items()
|
||||
if mcs.need_override_by_subclass(attr_name, attr)
|
||||
}
|
||||
missing = methods.difference(overridden)
|
||||
if missing:
|
||||
raise TypeError(
|
||||
f"Found missing callback method: {missing} "
|
||||
f"in class {cls.__module__}.{cls.__qualname__}."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def need_override_by_subclass(mcs, attr_name: str, attr: Any) -> bool:
|
||||
return (
|
||||
(
|
||||
attr_name.startswith("on_")
|
||||
and not attr_name.startswith("on_trainer_init")
|
||||
)
|
||||
or attr_name == "setup"
|
||||
) and callable(attr)
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class Callback(metaclass=_CallbackMeta):
|
||||
"""Tune base callback that can be extended and passed to a ``TrialRunner``
|
||||
|
||||
Tune callbacks are called from within the ``TrialRunner`` class. There are
|
||||
several hooks that can be used, all of which are found in the submethod
|
||||
definitions of this base class.
|
||||
|
||||
The parameters passed to the ``**info`` dict vary between hooks. The
|
||||
parameters passed are described in the docstrings of the methods.
|
||||
|
||||
This example will print a metric each time a result is received:
|
||||
|
||||
.. testcode::
|
||||
|
||||
from ray import tune
|
||||
from ray.tune import Callback
|
||||
|
||||
|
||||
class MyCallback(Callback):
|
||||
def on_trial_result(self, iteration, trials, trial, result,
|
||||
**info):
|
||||
print(f"Got result: {result['metric']}")
|
||||
|
||||
|
||||
def train_func(config):
|
||||
for i in range(10):
|
||||
tune.report(metric=i)
|
||||
|
||||
tuner = tune.Tuner(
|
||||
train_func,
|
||||
run_config=tune.RunConfig(
|
||||
callbacks=[MyCallback()]
|
||||
)
|
||||
)
|
||||
tuner.fit()
|
||||
|
||||
.. testoutput::
|
||||
:hide:
|
||||
|
||||
...
|
||||
"""
|
||||
|
||||
# File templates for any artifacts written by this callback
|
||||
# These files should live in the `trial.local_path` for each trial.
|
||||
# TODO(ml-team): Make this more visible to users to override. Internal use for now.
|
||||
_SAVED_FILE_TEMPLATES = []
|
||||
|
||||
# arguments here match Experiment.public_spec
|
||||
def setup(
|
||||
self,
|
||||
stop: Optional["Stopper"] = None,
|
||||
num_samples: Optional[int] = None,
|
||||
total_num_samples: Optional[int] = None,
|
||||
**info,
|
||||
):
|
||||
"""Called once at the very beginning of training.
|
||||
|
||||
Any Callback setup should be added here (setting environment
|
||||
variables, etc.)
|
||||
|
||||
Arguments:
|
||||
stop: Stopping criteria.
|
||||
If ``time_budget_s`` was passed to ``tune.RunConfig``, a
|
||||
``TimeoutStopper`` will be passed here, either by itself
|
||||
or as a part of a ``CombinedStopper``.
|
||||
num_samples: Number of times to sample from the
|
||||
hyperparameter space. Defaults to 1. If `grid_search` is
|
||||
provided as an argument, the grid will be repeated
|
||||
`num_samples` of times. If this is -1, (virtually) infinite
|
||||
samples are generated until a stopping condition is met.
|
||||
total_num_samples: Total number of samples factoring
|
||||
in grid search samplers.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_step_begin(self, iteration: int, trials: List["Trial"], **info):
|
||||
"""Called at the start of each tuning loop step.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_step_end(self, iteration: int, trials: List["Trial"], **info):
|
||||
"""Called at the end of each tuning loop step.
|
||||
|
||||
The iteration counter is increased before this hook is called.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_trial_start(
|
||||
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
||||
):
|
||||
"""Called after starting a trial instance.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
trial: Trial that just has been started.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_trial_restore(
|
||||
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
||||
):
|
||||
"""Called after restoring a trial instance.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
trial: Trial that just has been restored.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_trial_save(
|
||||
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
||||
):
|
||||
"""Called after receiving a checkpoint from a trial.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
trial: Trial that just saved a checkpoint.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_trial_result(
|
||||
self,
|
||||
iteration: int,
|
||||
trials: List["Trial"],
|
||||
trial: "Trial",
|
||||
result: Dict,
|
||||
**info,
|
||||
):
|
||||
"""Called after receiving a result from a trial.
|
||||
|
||||
The search algorithm and scheduler are notified before this
|
||||
hook is called.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
trial: Trial that just sent a result.
|
||||
result: Result that the trial sent.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_trial_complete(
|
||||
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
||||
):
|
||||
"""Called after a trial instance completed.
|
||||
|
||||
The search algorithm and scheduler are notified before this
|
||||
hook is called.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
trial: Trial that just has been completed.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_trial_recover(
|
||||
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
||||
):
|
||||
"""Called after a trial instance failed (errored) but the trial is scheduled
|
||||
for retry.
|
||||
|
||||
The search algorithm and scheduler are not notified.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
trial: Trial that just has errored.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_trial_error(
|
||||
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
||||
):
|
||||
"""Called after a trial instance failed (errored).
|
||||
|
||||
The search algorithm and scheduler are notified before this
|
||||
hook is called.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
trial: Trial that just has errored.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_checkpoint(
|
||||
self,
|
||||
iteration: int,
|
||||
trials: List["Trial"],
|
||||
trial: "Trial",
|
||||
checkpoint: "ray.tune.Checkpoint",
|
||||
**info,
|
||||
):
|
||||
"""Called after a trial saved a checkpoint with Tune.
|
||||
|
||||
Arguments:
|
||||
iteration: Number of iterations of the tuning loop.
|
||||
trials: List of trials.
|
||||
trial: Trial that just has errored.
|
||||
checkpoint: Checkpoint object that has been saved
|
||||
by the trial.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_experiment_end(self, trials: List["Trial"], **info):
|
||||
"""Called after experiment is over and all trials have concluded.
|
||||
|
||||
Arguments:
|
||||
trials: List of trials.
|
||||
**info: Kwargs dict for forward compatibility.
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_state(self) -> Optional[Dict]:
|
||||
"""Get the state of the callback.
|
||||
|
||||
This method should be implemented by subclasses to return a dictionary
|
||||
representation of the object's current state.
|
||||
|
||||
This is called automatically by Tune to periodically checkpoint callback state.
|
||||
Upon :ref:`Tune experiment restoration <tune-experiment-level-fault-tolerance>`,
|
||||
callback state will be restored via :meth:`~ray.tune.Callback.set_state`.
|
||||
|
||||
.. testcode::
|
||||
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from ray.tune import Callback
|
||||
from ray.tune.experiment import Trial
|
||||
|
||||
class MyCallback(Callback):
|
||||
def __init__(self):
|
||||
self._trial_ids = set()
|
||||
|
||||
def on_trial_start(
|
||||
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
||||
):
|
||||
self._trial_ids.add(trial.trial_id)
|
||||
|
||||
def get_state(self) -> Optional[Dict]:
|
||||
return {"trial_ids": self._trial_ids.copy()}
|
||||
|
||||
def set_state(self, state: Dict) -> Optional[Dict]:
|
||||
self._trial_ids = state["trial_ids"]
|
||||
|
||||
Returns:
|
||||
dict: State of the callback. Should be `None` if the callback does not
|
||||
have any state to save (this is the default).
|
||||
"""
|
||||
return None
|
||||
|
||||
def set_state(self, state: Dict):
|
||||
"""Set the state of the callback.
|
||||
|
||||
This method should be implemented by subclasses to restore the callback's
|
||||
state based on the given dict state.
|
||||
|
||||
This is used automatically by Tune to restore checkpoint callback state
|
||||
on :ref:`Tune experiment restoration <tune-experiment-level-fault-tolerance>`.
|
||||
|
||||
See :meth:`~ray.tune.Callback.get_state` for an example implementation.
|
||||
|
||||
Args:
|
||||
state: State of the callback.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class CallbackList(Callback):
|
||||
"""Call multiple callbacks at once."""
|
||||
|
||||
IS_CALLBACK_CONTAINER = True
|
||||
CKPT_FILE_TMPL = "callback-states-{}.pkl"
|
||||
|
||||
def __init__(self, callbacks: List[Callback]):
|
||||
self._callbacks = callbacks
|
||||
|
||||
def setup(self, **info):
|
||||
for callback in self._callbacks:
|
||||
try:
|
||||
callback.setup(**info)
|
||||
except TypeError as e:
|
||||
if "argument" in str(e):
|
||||
warnings.warn(
|
||||
"Please update `setup` method in callback "
|
||||
f"`{callback.__class__}` to match the method signature"
|
||||
" in `ray.tune.callback.Callback`.",
|
||||
FutureWarning,
|
||||
)
|
||||
callback.setup()
|
||||
else:
|
||||
raise e
|
||||
|
||||
def on_step_begin(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_step_begin(**info)
|
||||
|
||||
def on_step_end(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_step_end(**info)
|
||||
|
||||
def on_trial_start(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_trial_start(**info)
|
||||
|
||||
def on_trial_restore(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_trial_restore(**info)
|
||||
|
||||
def on_trial_save(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_trial_save(**info)
|
||||
|
||||
def on_trial_result(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_trial_result(**info)
|
||||
|
||||
def on_trial_complete(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_trial_complete(**info)
|
||||
|
||||
def on_trial_recover(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_trial_recover(**info)
|
||||
|
||||
def on_trial_error(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_trial_error(**info)
|
||||
|
||||
def on_checkpoint(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_checkpoint(**info)
|
||||
|
||||
def on_experiment_end(self, **info):
|
||||
for callback in self._callbacks:
|
||||
callback.on_experiment_end(**info)
|
||||
|
||||
def get_state(self) -> Optional[Dict]:
|
||||
"""Gets the state of all callbacks contained within this list.
|
||||
If there are no stateful callbacks, then None will be returned in order
|
||||
to avoid saving an unnecessary callback checkpoint file."""
|
||||
state = {}
|
||||
any_stateful_callbacks = False
|
||||
for i, callback in enumerate(self._callbacks):
|
||||
callback_state = callback.get_state()
|
||||
if callback_state:
|
||||
any_stateful_callbacks = True
|
||||
state[i] = callback_state
|
||||
if not any_stateful_callbacks:
|
||||
return None
|
||||
return state
|
||||
|
||||
def set_state(self, state: Dict):
|
||||
"""Sets the state for all callbacks contained within this list.
|
||||
Skips setting state for all stateless callbacks where `get_state`
|
||||
returned None."""
|
||||
for i, callback in enumerate(self._callbacks):
|
||||
callback_state = state.get(i, None)
|
||||
if callback_state:
|
||||
callback.set_state(callback_state)
|
||||
|
||||
def save_to_dir(self, checkpoint_dir: str, session_str: str = "default"):
|
||||
"""Save the state of the callback list to the checkpoint_dir.
|
||||
|
||||
Args:
|
||||
checkpoint_dir: directory where the checkpoint is stored.
|
||||
session_str: Unique identifier of the current run session (ex: timestamp).
|
||||
"""
|
||||
state_dict = self.get_state()
|
||||
|
||||
if state_dict:
|
||||
file_name = self.CKPT_FILE_TMPL.format(session_str)
|
||||
tmp_file_name = f"tmp-{file_name}"
|
||||
_atomic_save(
|
||||
state=state_dict,
|
||||
checkpoint_dir=checkpoint_dir,
|
||||
file_name=file_name,
|
||||
tmp_file_name=tmp_file_name,
|
||||
)
|
||||
|
||||
def restore_from_dir(self, checkpoint_dir: str):
|
||||
"""Restore the state of the list of callbacks from the checkpoint_dir.
|
||||
|
||||
You should check if it's possible to restore with `can_restore`
|
||||
before calling this method.
|
||||
|
||||
Args:
|
||||
checkpoint_dir: directory where the checkpoint is stored.
|
||||
|
||||
Raises:
|
||||
RuntimeError: if unable to find checkpoint.
|
||||
NotImplementedError: if the `set_state` method is not implemented.
|
||||
"""
|
||||
state_dict = _load_newest_checkpoint(
|
||||
checkpoint_dir, self.CKPT_FILE_TMPL.format("*")
|
||||
)
|
||||
if not state_dict:
|
||||
raise RuntimeError(
|
||||
"Unable to find checkpoint in {}.".format(checkpoint_dir)
|
||||
)
|
||||
self.set_state(state_dict)
|
||||
|
||||
def can_restore(self, checkpoint_dir: str) -> bool:
|
||||
"""Check if the checkpoint_dir contains the saved state for this callback list.
|
||||
|
||||
Args:
|
||||
checkpoint_dir: Directory to look for a saved state file in.
|
||||
|
||||
Returns:
|
||||
can_restore: True if the checkpoint_dir contains a file of the
|
||||
format `CKPT_FILE_TMPL`. False otherwise.
|
||||
"""
|
||||
return any(
|
||||
glob.iglob(Path(checkpoint_dir, self.CKPT_FILE_TMPL.format("*")).as_posix())
|
||||
)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._callbacks)
|
||||
|
||||
def __getitem__(self, i: int) -> "Callback":
|
||||
return self._callbacks[i]
|
||||
Reference in New Issue
Block a user