Repository: ray-project/ray

Issue:
[tune] New feature: callback support for Tune
<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->

### Describe your feature request

Callbacks are a very useful addition to tuning libraries, allowing to eg. measure progress of long tuning processes. Adding the possibility to use them in Ray Tune, with, for example:

```python
tune.run(
	callbacks=[func, func2]
)
```

would be very helpful and bring it in line with other hyperparameter tuning libraries like Optuna.


Discussion:
@richardliaw 
cc @ray-project/ray-tune 
Hi @Yard1, thanks for the issue. We're looking into supporting callbacks very soon. The main question I have is that there a many possibilities to inject callbacks. It would be great if you could provide some insight into your use case so we can make sure that we cover it with our implementation.

For instance, callbacks could be called every time when a result is received from a trial, i.e. after calling `tune.report()`. The callback would then receive the reference to the trial, and then has access to the last result and other properties. Would this be useful for you, or do you need callback support at other places or with other data?

Another possibility would be to support callbacks on the outer training loop of the trial runner, i.e. before and after trial events for all trials are executed. 

 It would be great to learn what you are trying to achieve with the callbacks. Thanks!
For me, the most important thing is being able to measure progress. I would like to be able to tell how many trials have been completed out of total use that information by eg. incrementing a counter in a progress bar. I think that the simplest way would be to have a callback at the end of a trial, so after `tune.report()` sounds like something I want.

I imagine that some people would want callbacks before a trial starts, to eg. log parameters, but this isn't something I need myself. 

Resolving patch:
diff --git a/python/ray/tune/__init__.py b/python/ray/tune/__init__.py
--- a/python/ray/tune/__init__.py
+++ b/python/ray/tune/__init__.py
@@ -8,6 +8,7 @@
 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.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/trial_runner.py b/python/ray/tune/trial_runner.py
--- 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 @@ def _from_cloudpickle(self, obj):
         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 @@ def __init__(self,
                  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 @@ def __init__(self,
                 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 @@ def step(self):
             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 @@ def step(self):
                 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 @@ def _process_events(self):
             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 @@ def _process_trial(self, trial):
                 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 @@ def _process_trial(self, trial):
                 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 @@ def _process_trial_failure(self, trial, error_msg):
             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 @@ def __getstate__(self):
         """
         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
--- 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:


Test patch:
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
--- /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__]))