[Tune] Save and restore stateful callbacks as part of experiment checkpoint

python/ray/tune/BUILD

@@ -39,6 +39,14 @@ py_test(
     tags = ["team:ml", "exclusive"],
 )
 
+py_test(
+    name = "test_callbacks",
+    size = "small",
+    srcs = ["tests/test_callbacks.py"],
+    deps = [":tune_lib"],
+    tags = ["team:ml", "exclusive"],
+)
+
 py_test(
     name = "test_checkpoint_manager",
     size = "small",

python/ray/tune/callback.py

@@ -1,8 +1,11 @@
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 from abc import ABCMeta
+import glob
+import os
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 import warnings
 
 from ray.util.annotations import PublicAPI, DeveloperAPI
+from ray.tune.utils.util import _atomic_save, _load_newest_checkpoint
 
 if TYPE_CHECKING:
     from ray.air._internal.checkpoint_manager import _TrackedCheckpoint
@@ -104,6 +107,8 @@ def train(config):
 
     """
 
+    CKPT_FILE_TMPL = "callback-state-{}.pkl"
+
     # arguments here match Experiment.public_spec
     def setup(
         self,
@@ -278,12 +283,91 @@ def on_experiment_end(self, trials: List["Trial"], **info):
         """
         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.
+
+        Returns:
+            state: 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):
+        """Get the state of the callback.
+
+        This method should be implemented by subclasses to restore the callback's
+        state based on the given dict state.
+
+        Args:
+            state: State of the callback.
+        """
+        pass
+
+    def can_restore(self, checkpoint_dir: str) -> bool:
+        """Check if the checkpoint_dir contains the saved state for this callback.
+
+        Returns:
+            can_restore: True if the checkpoint_dir contains a file of the
+                format `CKPT_FILE_TMPL`. False otherwise.
+        """
+        return bool(
+            glob.glob(os.path.join(checkpoint_dir, self.CKPT_FILE_TMPL.format("*")))
+        )
+
+    def save_to_dir(self, checkpoint_dir: str, session_str: str = "default"):
+        """Save the state of the callback to the checkpoint_dir.
+
+        Args:
+            checkpoint_dir: directory where the checkpoint is stored.
+            session_str: Unique identifier of the current run session (ex: timestamp).
+
+        Raises:
+            NotImplementedError: if the `get_state` method is not implemented.
+        """
+        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 object 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)
+
 
 @DeveloperAPI
 class CallbackList(Callback):
     """Call multiple callbacks at once."""
 
     IS_CALLBACK_CONTAINER = True
+    CKPT_FILE_TMPL = "callback-list-state-{}.pkl"
 
     def __init__(self, callbacks: List[Callback]):
         self._callbacks = callbacks
@@ -343,3 +427,21 @@ def on_checkpoint(self, **info):
     def on_experiment_end(self, **info):
         for callback in self._callbacks:
             callback.on_experiment_end(**info)
+
+    def get_state(self) -> Optional[Dict]:
+        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):
+        for i, callback in enumerate(self._callbacks):
+            callback_state = state.get(i, None)
+            if callback_state:
+                callback.set_state(callback_state)

python/ray/tune/execution/trial_runner.py

@@ -151,6 +151,7 @@ def checkpoint(
         trial_runner: "TrialRunner",
         trial_executor: RayTrialExecutor,
         search_alg: SearchAlgorithm,
+        callbacks: CallbackList,
         force: bool = False,
     ):
         """Saves execution state to `self._local_checkpoint_dir`.
@@ -190,6 +191,9 @@ def _serialize_and_write():
             search_alg.save_to_dir(
                 self._local_checkpoint_dir, session_str=self._session_str
             )
+            callbacks.save_to_dir(
+                self._local_checkpoint_dir, session_str=self._session_str
+            )
 
         checkpoint_time_start = time.monotonic()
         with out_of_band_serialize_dataset():
@@ -347,6 +351,7 @@ def __init__(
         self._search_alg = search_alg or BasicVariantGenerator()
         self._scheduler_alg = scheduler or FIFOScheduler()
         self.trial_executor = trial_executor or RayTrialExecutor()
+        self._callbacks = CallbackList(callbacks or [])
         self._insufficient_resources_manager = _InsufficientResourcesManager()
         self._pending_trial_queue_times = {}
 
@@ -469,8 +474,6 @@ def __init__(
                 TrialRunner.CKPT_FILE_TMPL.format(self._session_str),
             )
 
-        self._callbacks = CallbackList(callbacks or [])
-
         if checkpoint_period is None:
             checkpoint_period = os.getenv("TUNE_GLOBAL_CHECKPOINT_S", "auto")
 
@@ -751,6 +754,7 @@ def checkpoint(self, force: bool = False):
                 trial_runner=self,
                 trial_executor=self.trial_executor,
                 search_alg=self._search_alg,
+                callbacks=self._callbacks,
                 force=force,
             )
 
@@ -795,10 +799,13 @@ def resume(
         # 1. Restore trial runner state
         self.__setstate__(runner_state["runner_data"])
 
-        # 2. Restore search algorithm state
+        # 2. Restore search algorithm and callback state
         if self._search_alg.has_checkpoint(self._local_checkpoint_dir):
             self._search_alg.restore_from_dir(self._local_checkpoint_dir)
 
+        if self._callbacks.can_restore(self._local_checkpoint_dir):
+            self._callbacks.restore_from_dir(self._local_checkpoint_dir)
+
         # 3. Load trial table from experiment checkpoint
         trials = []
         for trial_json_state in runner_state["checkpoints"]:

python/ray/tune/tests/test_callbacks.py

@@ -0,0 +1,92 @@
+import pytest
+from typing import Dict, Optional
+
+from ray.tune.callback import Callback, CallbackList
+
+
+class StatefulCallback(Callback):
+    CKPT_FILE_TMPL = "test-callback-state-{}.json"
+
+    def __init__(self):
+        self.counter = 0
+
+    def on_trial_result(self, iteration, trials, trial, result, **info):
+        self.counter += 1
+
+    def get_state(self) -> Optional[Dict]:
+        return {"counter": self.counter}
+
+    def set_state(self, state: Dict):
+        self.counter = state["counter"]
+
+
+def test_stateful_callback_save_and_restore(tmp_path):
+    """Checks that a stateful callback can be saved to a directory and restored with
+    the right state."""
+
+    callback = StatefulCallback()
+    for i in range(3):
+        callback.on_trial_result(i, None, None, None)
+    callback.save_to_dir(str(tmp_path))
+    assert list(tmp_path.glob(StatefulCallback.CKPT_FILE_TMPL.format("*")))
+    assert callback.can_restore(str(tmp_path))
+
+    restored_callback = StatefulCallback()
+    restored_callback.restore_from_dir(str(tmp_path))
+    assert restored_callback.counter == 3
+
+
+def test_stateless_callback_save_and_restore(tmp_path):
+    """Checks that proper errors are raised/handled when saving/restoring a
+    stateless callback (i.e. one that doesn't implement get/set_state)."""
+
+    class StatelessCallback(Callback):
+        def handle_save_error(self, error: Exception):
+            assert isinstance(error, NotImplementedError)
+
+    callback = StatelessCallback()
+    callback.save_to_dir(str(tmp_path))
+
+    assert not list(tmp_path.glob(StatelessCallback.CKPT_FILE_TMPL.format("*")))
+    assert not callback.can_restore(str(tmp_path))
+    with pytest.raises(RuntimeError):
+        callback.restore_from_dir(str(tmp_path))
+
+
+def test_callback_list_with_stateful_callback(tmp_path):
+    """Checks that a callback list saves and restores all callbacks contained
+    inside it."""
+
+    callbacks = CallbackList([Callback(), StatefulCallback()])
+    for i in range(3):
+        callbacks.on_trial_result(iteration=i, trials=None, trial=None, result=None)
+
+    callbacks.save_to_dir(str(tmp_path))
+
+    assert list(tmp_path.glob(CallbackList.CKPT_FILE_TMPL.format("*")))
+    assert callbacks.can_restore(str(tmp_path))
+
+    restored_callbacks = CallbackList([Callback(), StatefulCallback()])
+    restored_callbacks.restore_from_dir(str(tmp_path))
+
+    assert restored_callbacks._callbacks[1].counter == 3
+
+
+def test_callback_list_without_stateful_callback(tmp_path):
+    """If no callbacks within a CallbackList are stateful, then nothing
+    should be saved."""
+
+    callbacks = CallbackList([Callback(), Callback()])
+    callbacks.save_to_dir(str(tmp_path))
+
+    assert not list(tmp_path.glob(CallbackList.CKPT_FILE_TMPL.format("*")))
+    assert not callbacks.can_restore(str(tmp_path))
+
+    with pytest.raises(RuntimeError):
+        callbacks.restore_from_dir(str(tmp_path))
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_trial_runner_3.py

@@ -32,6 +32,7 @@
 from ray.tune.search.search_generator import SearchGenerator
 from ray.tune.syncer import SyncConfig, Syncer
 from ray.tune.tests.tune_test_util import TrialResultObserver
+from ray.tune.tests.test_callbacks import StatefulCallback
 
 
 class MyCallbacks(DefaultCallbacks):
@@ -402,6 +403,29 @@ def num_running_trials():
         count = Counter(evaluated)
         assert all(v <= 3 for v in count.values())
 
+    def testCallbackSaveRestore(self):
+        """Check that experiment state save + restore handles stateful callbacks."""
+        ray.init(num_cpus=2)
+        runner = TrialRunner(
+            local_checkpoint_dir=self.tmpdir,
+            callbacks=[StatefulCallback()],
+            trial_executor=RayTrialExecutor(resource_manager=self._resourceManager()),
+        )
+        runner.add_trial(Trial("__fake", stub=True))
+        for i in range(3):
+            runner._callbacks.on_trial_result(
+                iteration=i, trials=None, trial=None, result=None
+            )
+        runner.checkpoint(force=True)
+        callback = StatefulCallback()
+        runner2 = TrialRunner(
+            local_checkpoint_dir=self.tmpdir,
+            callbacks=[callback],
+        )
+        assert callback.counter == 0
+        runner2.resume()
+        assert callback.counter == 3
+
     def testTrialErrorResumeFalse(self):
         ray.init(num_cpus=3, local_mode=True, include_dashboard=False)
         runner = TrialRunner(