chore: import upstream snapshot with attribution

python/ray/tune/tests/test_callbacks.py

@@ -0,0 +1,60 @@
+from typing import Dict, Optional
+
+import pytest
+
+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_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__]))