[tune] Callbacks for tune runs (#11001)

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",

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",

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__]))

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)

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: