chore: import upstream snapshot with attribution

python/ray/tune/tests/execution/test_controller_callback_integration.py

@@ -0,0 +1,69 @@
+import sys
+from typing import Dict, Optional
+
+import pytest
+
+import ray
+from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
+from ray.train.tests.util import mock_storage_context
+from ray.tune import Callback, ResumeConfig
+from ray.tune.execution.tune_controller import TuneController
+from ray.tune.experiment import Trial
+from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
+
+
+@pytest.fixture(scope="function")
+def ray_start_4_cpus_2_gpus_extra():
+    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def register_test_trainable():
+    register_mock_trainable()
+
+
+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"]
+
+
+@pytest.mark.parametrize(
+    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
+)
+def test_callback_save_restore(
+    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir
+):
+    """Check that callback state is restored correctly.
+
+    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCallbackSaveRestore
+    """
+    storage = mock_storage_context()
+    runner = TuneController(callbacks=[StatefulCallback()], storage=storage)
+    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, stub=True, storage=storage))
+    for i in range(3):
+        runner._callbacks.on_trial_result(
+            iteration=i, trials=None, trial=None, result=None
+        )
+    runner.checkpoint(force=True, wait=True)
+    callback = StatefulCallback()
+    runner2 = TuneController(callbacks=[callback], storage=storage)
+    assert callback.counter == 0
+    runner2.resume(resume_config=ResumeConfig())
+    assert callback.counter == 3
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))