diff --git a/doc/source/tune/api_docs/internals.rst b/doc/source/tune/api_docs/internals.rst index 644a0f42b..1ac7b4852 100644 --- a/doc/source/tune/api_docs/internals.rst +++ b/doc/source/tune/api_docs/internals.rst @@ -111,6 +111,15 @@ Trial .. autoclass:: ray.tune.trial.Trial +.. _tune-callbacks-docs: + +Callbacks +--------- + +.. autoclass:: ray.tune.trial_runner.Callback + :members: + + .. _resources-docstring: Resources diff --git a/doc/source/tune/user-guide.rst b/doc/source/tune/user-guide.rst index 43facb13f..0a9524a16 100644 --- a/doc/source/tune/user-guide.rst +++ b/doc/source/tune/user-guide.rst @@ -509,6 +509,40 @@ too. If ``log_to_file`` is set, Tune will automatically register a new logging handler for Ray's base logger and log the output to the specified stderr output file. +.. _tune-callbacks: + +Callbacks +--------- + +Ray Tune supports callbacks that are called during various times of the training process. +Callbacks can be passed as a parameter to ``tune.run()``, and the submethod will be +invoked automatically. + +This simple callback just prints a metric each time a result is received: + +.. code-block:: python + + 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(config): + for i in range(10): + tune.report(metric=i) + + + tune.run( + train, + callbacks=[MyCallback()]) + +For more details and available hooks, please :ref:`see the API docs for Ray Tune callbacks `. + + .. _tune-debugging: Debugging diff --git a/python/ray/tune/BUILD b/python/ray/tune/BUILD index f85583f74..44034097e 100644 --- a/python/ray/tune/BUILD +++ b/python/ray/tune/BUILD @@ -201,6 +201,14 @@ py_test( tags = ["exclusive"], ) +py_test( + name = "test_trial_runner_callbacks", + size = "small", + srcs = ["tests/test_trial_runner_callbacks.py"], + deps = [":tune_lib"], + tags = ["exclusive"], +) + py_test( name = "test_var", size = "small", diff --git a/python/ray/tune/__init__.py b/python/ray/tune/__init__.py index 813b05760..f3babecbb 100644 --- a/python/ray/tune/__init__.py +++ b/python/ray/tune/__init__.py @@ -8,6 +8,7 @@ from ray.tune.stopper import Stopper, EarlyStopping from ray.tune.registry import register_env, register_trainable from ray.tune.trainable import Trainable from ray.tune.durable_trainable import DurableTrainable +from ray.tune.trial_runner import Callback from ray.tune.suggest import grid_search from ray.tune.session import ( report, get_trial_dir, get_trial_name, get_trial_id, make_checkpoint_dir, @@ -21,7 +22,7 @@ from ray.tune.suggest import create_searcher from ray.tune.schedulers import create_scheduler __all__ = [ - "Trainable", "DurableTrainable", "TuneError", "grid_search", + "Trainable", "DurableTrainable", "TuneError", "Callback", "grid_search", "register_env", "register_trainable", "run", "run_experiments", "with_parameters", "Stopper", "EarlyStopping", "Experiment", "function", "sample_from", "track", "uniform", "quniform", "choice", "randint", diff --git a/python/ray/tune/tests/test_trial_runner_callbacks.py b/python/ray/tune/tests/test_trial_runner_callbacks.py new file mode 100644 index 000000000..13deb6e7b --- /dev/null +++ b/python/ray/tune/tests/test_trial_runner_callbacks.py @@ -0,0 +1,202 @@ +import os +import shutil +import sys +import tempfile +import time +import unittest + +import ray +from ray import tune +from ray.rllib import _register_all +from ray.tune.checkpoint_manager import Checkpoint +from ray.tune.ray_trial_executor import RayTrialExecutor +from ray.tune.result import TRAINING_ITERATION + +from ray.tune.trial import Trial +from ray.tune.trial_runner import Callback, TrialRunner + + +class TestCallback(Callback): + def __init__(self): + self.state = {} + + def on_step_begin(self, **info): + self.state["step_begin"] = info + + def on_step_end(self, **info): + self.state["step_end"] = info + + def on_trial_start(self, **info): + self.state["trial_start"] = info + + def on_trial_restore(self, **info): + self.state["trial_restore"] = info + + def on_trial_save(self, **info): + self.state["trial_save"] = info + + def on_trial_result(self, **info): + self.state["trial_result"] = info + result = info["result"] + trial = info["trial"] + assert result.get(TRAINING_ITERATION, None) != trial.last_result.get( + TRAINING_ITERATION, None) + + def on_trial_complete(self, **info): + self.state["trial_complete"] = info + + def on_trial_fail(self, **info): + self.state["trial_fail"] = info + + +class _MockTrialExecutor(RayTrialExecutor): + def __init__(self): + super().__init__() + self.results = {} + self.next_trial = None + self.failed_trial = None + + def fetch_result(self, trial): + return self.results.get(trial, {}) + + def get_next_available_trial(self): + return self.next_trial or super().get_next_available_trial() + + def get_next_failed_trial(self): + return self.failed_trial or super().get_next_failed_trial() + + +class TrialRunnerCallbacks(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.callback = TestCallback() + self.executor = _MockTrialExecutor() + self.trial_runner = TrialRunner( + trial_executor=self.executor, callbacks=[self.callback]) + + def tearDown(self): + ray.shutdown() + _register_all() # re-register the evicted objects + if "CUDA_VISIBLE_DEVICES" in os.environ: + del os.environ["CUDA_VISIBLE_DEVICES"] + shutil.rmtree(self.tmpdir) + + def testCallbackSteps(self): + trials = [ + Trial("__fake", trial_id="one"), + Trial("__fake", trial_id="two") + ] + for t in trials: + self.trial_runner.add_trial(t) + + self.trial_runner.step() + + # Trial 1 has been started + self.assertEqual(self.callback.state["trial_start"]["iteration"], 0) + self.assertEqual(self.callback.state["trial_start"]["trial"].trial_id, + "one") + + # All these events haven't happened, yet + self.assertTrue( + all(k not in self.callback.state for k in [ + "trial_restore", "trial_save", "trial_result", + "trial_complete", "trial_fail" + ])) + + self.trial_runner.step() + + # Iteration not increased yet + self.assertEqual(self.callback.state["step_begin"]["iteration"], 1) + + # Iteration increased + self.assertEqual(self.callback.state["step_end"]["iteration"], 2) + + # Second trial has been just started + self.assertEqual(self.callback.state["trial_start"]["iteration"], 1) + self.assertEqual(self.callback.state["trial_start"]["trial"].trial_id, + "two") + + cp = Checkpoint(Checkpoint.PERSISTENT, "__checkpoint", + {TRAINING_ITERATION: 0}) + + # Let the first trial save a checkpoint + trials[0].saving_to = cp + self.trial_runner.step() + self.assertEqual(self.callback.state["trial_save"]["iteration"], 2) + self.assertEqual(self.callback.state["trial_save"]["trial"].trial_id, + "one") + + # Let the second trial send a result + result = {TRAINING_ITERATION: 1, "metric": 800, "done": False} + self.executor.results[trials[1]] = result + self.executor.next_trial = trials[1] + self.assertEqual(trials[1].last_result, {}) + self.trial_runner.step() + self.assertEqual(self.callback.state["trial_result"]["iteration"], 3) + self.assertEqual(self.callback.state["trial_result"]["trial"].trial_id, + "two") + self.assertEqual( + self.callback.state["trial_result"]["result"]["metric"], 800) + self.assertEqual(trials[1].last_result["metric"], 800) + + # Let the second trial restore from a checkpoint + trials[1].restoring_from = cp + self.executor.results[trials[1]] = trials[1].last_result + self.trial_runner.step() + self.assertEqual(self.callback.state["trial_restore"]["iteration"], 4) + self.assertEqual( + self.callback.state["trial_restore"]["trial"].trial_id, "two") + + # Let the second trial finish + trials[1].restoring_from = None + self.executor.results[trials[1]] = { + TRAINING_ITERATION: 2, + "metric": 900, + "done": True + } + self.trial_runner.step() + self.assertEqual(self.callback.state["trial_complete"]["iteration"], 5) + self.assertEqual( + self.callback.state["trial_complete"]["trial"].trial_id, "two") + + # Let the first trial error + self.executor.failed_trial = trials[0] + self.trial_runner.step() + self.assertEqual(self.callback.state["trial_fail"]["iteration"], 6) + self.assertEqual(self.callback.state["trial_fail"]["trial"].trial_id, + "one") + + def testCallbacksEndToEnd(self): + def train(config): + if config["do"] == "save": + with tune.checkpoint_dir(0): + pass + tune.report(metric=1) + elif config["do"] == "fail": + raise RuntimeError("I am failing on purpose.") + elif config["do"] == "delay": + time.sleep(2) + tune.report(metric=20) + + config = {"do": tune.grid_search(["save", "fail", "delay"])} + + tune.run( + train, + config=config, + raise_on_failed_trial=False, + callbacks=[self.callback]) + + self.assertEqual( + self.callback.state["trial_fail"]["trial"].config["do"], "fail") + self.assertEqual( + self.callback.state["trial_save"]["trial"].config["do"], "save") + self.assertEqual( + self.callback.state["trial_result"]["trial"].config["do"], "delay") + self.assertEqual( + self.callback.state["trial_complete"]["trial"].config["do"], + "delay") + + +if __name__ == "__main__": + import pytest + sys.exit(pytest.main(["-v", __file__])) diff --git a/python/ray/tune/trial_runner.py b/python/ray/tune/trial_runner.py index 34acc47da..8cb8f2b83 100644 --- a/python/ray/tune/trial_runner.py +++ b/python/ray/tune/trial_runner.py @@ -1,3 +1,5 @@ +from typing import Dict, List + import click from datetime import datetime import json @@ -69,6 +71,186 @@ class _TuneFunctionDecoder(json.JSONDecoder): return cloudpickle.loads(hex_to_binary(obj["value"])) +class Callback: + """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: + + .. code-block:: python + + 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(config): + for i in range(10): + tune.report(metric=i) + + + tune.run( + train, + callbacks=[MyCallback()]) + + """ + + def on_step_begin(self, iteration: int, trials: List[Trial], **info): + """Called at the start of each tuning loop step. + + Arguments: + iteration (int): Number of iterations of the tuning loop. + trials (List[Trial]): 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 (int): Number of iterations of the tuning loop. + trials (List[Trial]): 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 (int): Number of iterations of the tuning loop. + trials (List[Trial]): List of trials. + trial (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 (int): Number of iterations of the tuning loop. + trials (List[Trial]): List of trials. + trial (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 (int): Number of iterations of the tuning loop. + trials (List[Trial]): List of trials. + trial (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 (int): Number of iterations of the tuning loop. + trials (List[Trial]): List of trials. + trial (Trial): Trial that just sent a result. + result (Dict): 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 (int): Number of iterations of the tuning loop. + trials (List[Trial]): List of trials. + trial (Trial): Trial that just has been completed. + **info: Kwargs dict for forward compatibility. + """ + pass + + def on_trial_fail(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 (int): Number of iterations of the tuning loop. + trials (List[Trial]): List of trials. + trial (Trial): Trial that just has errored. + **info: Kwargs dict for forward compatibility. + """ + pass + + +class _CallbackList: + """Call multiple callbacks at once.""" + + def __init__(self, callbacks: List[Callback]): + self._callbacks = callbacks + + 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_fail(self, **info): + for callback in self._callbacks: + callback.on_trial_fail(**info) + + class TrialRunner: """A TrialRunner implements the event loop for scheduling trials on Ray. @@ -114,6 +296,9 @@ class TrialRunner: checkpoint_period (int): Trial runner checkpoint periodicity in seconds. Defaults to 10. trial_executor (TrialExecutor): Defaults to RayTrialExecutor. + callbacks (list): List of callbacks that will be called at different + times in the training loop. Must be instances of the + ``ray.tune.trial_runner.Callback`` class. """ CKPT_FILE_TMPL = "experiment_state-{}.json" @@ -133,6 +318,7 @@ class TrialRunner: verbose=True, checkpoint_period=None, trial_executor=None, + callbacks=None, metric=None): self._search_alg = search_alg or BasicVariantGenerator() self._scheduler_alg = scheduler or FIFOScheduler() @@ -214,6 +400,8 @@ class TrialRunner: self._local_checkpoint_dir, TrialRunner.CKPT_FILE_TMPL.format(self._session_str)) + self._callbacks = _CallbackList(callbacks or []) + @property def resumed(self): return self._resumed @@ -367,10 +555,17 @@ class TrialRunner: raise TuneError("Called step when all trials finished?") with warn_if_slow("on_step_begin"): self.trial_executor.on_step_begin(self) + with warn_if_slow("callbacks.on_step_begin"): + self._callbacks.on_step_begin( + iteration=self._iteration, trials=self._trials) next_trial = self._get_next_trial() # blocking if next_trial is not None: with warn_if_slow("start_trial"): self.trial_executor.start_trial(next_trial) + self._callbacks.on_trial_start( + iteration=self._iteration, + trials=self._trials, + trial=next_trial) elif self.trial_executor.get_running_trials(): self._process_events() # blocking else: @@ -393,6 +588,9 @@ class TrialRunner: self._server.shutdown() with warn_if_slow("on_step_end"): self.trial_executor.on_step_end(self) + with warn_if_slow("callbacks.on_step_end"): + self._callbacks.on_step_end( + iteration=self._iteration, trials=self._trials) def get_trial(self, tid): trial = [t for t in self._trials if t.trial_id == tid] @@ -479,9 +677,19 @@ class TrialRunner: if trial.is_restoring: with warn_if_slow("process_trial_restore"): self._process_trial_restore(trial) + with warn_if_slow("callbacks.on_trial_restore"): + self._callbacks.on_trial_restore( + iteration=self._iteration, + trials=self._trials, + trial=trial) elif trial.is_saving: with warn_if_slow("process_trial_save") as profile: self._process_trial_save(trial) + with warn_if_slow("callbacks.on_trial_save"): + self._callbacks.on_trial_save( + iteration=self._iteration, + trials=self._trials, + trial=trial) if profile.too_slow and trial.sync_on_checkpoint: # TODO(ujvl): Suggest using DurableTrainable once # API has converged. @@ -537,6 +745,10 @@ class TrialRunner: self._scheduler_alg.on_trial_complete(self, trial, flat_result) self._search_alg.on_trial_complete( trial.trial_id, result=flat_result) + self._callbacks.on_trial_complete( + iteration=self._iteration, + trials=self._trials, + trial=trial) decision = TrialScheduler.STOP else: with warn_if_slow("scheduler.on_trial_result"): @@ -545,10 +757,21 @@ class TrialRunner: with warn_if_slow("search_alg.on_trial_result"): self._search_alg.on_trial_result(trial.trial_id, flat_result) + with warn_if_slow("callbacks.on_trial_result"): + self._callbacks.on_trial_result( + iteration=self._iteration, + trials=self._trials, + trial=trial, + result=result.copy()) if decision == TrialScheduler.STOP: with warn_if_slow("search_alg.on_trial_complete"): self._search_alg.on_trial_complete( trial.trial_id, result=flat_result) + with warn_if_slow("callbacks.on_trial_complete"): + self._callbacks.on_trial_complete( + iteration=self._iteration, + trials=self._trials, + trial=trial) if not is_duplicate: trial.update_last_result( @@ -680,6 +903,10 @@ class TrialRunner: else: self._scheduler_alg.on_trial_error(self, trial) self._search_alg.on_trial_complete(trial.trial_id, error=True) + self._callbacks.on_trial_fail( + iteration=self._iteration, + trials=self._trials, + trial=trial) self.trial_executor.stop_trial( trial, error=True, error_msg=error_msg) @@ -845,13 +1072,8 @@ class TrialRunner: """ state = self.__dict__.copy() for k in [ - "_trials", - "_stop_queue", - "_server", - "_search_alg", - "_scheduler_alg", - "trial_executor", - "_syncer", + "_trials", "_stop_queue", "_server", "_search_alg", + "_scheduler_alg", "trial_executor", "_syncer", "_callbacks" ]: del state[k] state["launch_web_server"] = bool(self._server) diff --git a/python/ray/tune/tune.py b/python/ray/tune/tune.py index bc7d48a44..011b5520f 100644 --- a/python/ray/tune/tune.py +++ b/python/ray/tune/tune.py @@ -100,6 +100,7 @@ def run( reuse_actors=False, trial_executor=None, raise_on_failed_trial=True, + callbacks=None, # Deprecated args ray_auto_init=None, run_errored_only=None, @@ -259,6 +260,9 @@ def run( trial_executor (TrialExecutor): Manage the execution of trials. raise_on_failed_trial (bool): Raise TuneError if there exists failed trial (of ERROR state) when the experiments complete. + callbacks (list): List of callbacks that will be called at different + times in the training loop. Must be instances of the + ``ray.tune.trial_runner.Callback`` class. Returns: @@ -375,6 +379,7 @@ def run( verbose=bool(verbose > 1), fail_fast=fail_fast, trial_executor=trial_executor, + callbacks=callbacks, metric=metric) if not runner.resumed: