chore: import upstream snapshot with attribution

rllib/callbacks/tests/test_callbacks_on_algorithm.py

@@ -0,0 +1,115 @@
+import tempfile
+import time
+import unittest
+
+import ray
+from ray import tune
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.callbacks.callbacks import RLlibCallback
+from ray.rllib.examples.envs.classes.cartpole_crashing import CartPoleCrashing
+from ray.rllib.utils.test_utils import check
+
+
+def on_env_runners_recreated(
+    algorithm,
+    env_runner_group,
+    env_runner_indices,
+    is_evaluation,
+    **kwargs,
+):
+    # Store in the algorithm object's counters the number of times, this worker
+    # (ID'd by index and whether eval or not) has been recreated/restarted.
+    for id_ in env_runner_indices:
+        key = f"{'eval_' if is_evaluation else ''}worker_{id_}_recreated"
+        # Increase the counter.
+        algorithm.metrics.log_value(key, 1, reduce="lifetime_sum")
+        print(f"changed {key} to {algorithm._counters[key]}")
+
+    # Execute some dummy code on each of the recreated workers.
+    results = env_runner_group.foreach_env_runner(lambda w: w.ping())
+    print(results)  # should print "pong" n times (one for each recreated worker).
+
+
+class InitAndCheckpointRestoredCallbacks(RLlibCallback):
+    def on_algorithm_init(self, *, algorithm, metrics_logger, **kwargs):
+        self._on_init_was_called = True
+
+    def on_checkpoint_loaded(self, *, algorithm, **kwargs):
+        self._on_checkpoint_loaded_was_called = True
+
+
+class TestCallbacks(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_on_env_runners_recreated_callback(self):
+        tune.register_env("env", lambda cfg: CartPoleCrashing(cfg))
+
+        config = (
+            PPOConfig()
+            .environment("env", env_config={"p_crash": 1.0})
+            .callbacks(on_env_runners_recreated=on_env_runners_recreated)
+            .env_runners(num_env_runners=3)
+            .fault_tolerance(
+                restart_failed_env_runners=True,
+                delay_between_env_runner_restarts_s=0,
+            )
+        )
+
+        algo = config.build()
+        original_env_runner_ids = algo.env_runner_group.healthy_worker_ids()
+        for id_ in original_env_runner_ids:
+            check(algo.metrics.peek(f"worker_{id_}_recreated", default=0), 0)
+        check(algo.metrics.peek("total_num_workers_recreated", default=0), 0)
+
+        # After building the algorithm, we should have 2 healthy (remote) workers.
+        self.assertTrue(len(original_env_runner_ids) == 3)
+
+        # Train a bit (and have the envs/workers crash).
+        for _ in range(3):
+            print(algo.train())
+            time.sleep(15.0)
+
+        algo.restore_env_runners(algo.env_runner_group)
+        # After training, the `on_workers_recreated` callback should have captured
+        # the exact worker IDs recreated (the exact number of times) as the actor
+        # manager itself. This confirms that the callback is triggered correctly,
+        # always.
+        new_worker_ids = algo.env_runner_group.healthy_worker_ids()
+        self.assertEqual(len(new_worker_ids), 3)
+        for id_ in new_worker_ids:
+            # num_restored = algo.env_runner_group.restored_actors_history[id_]
+            self.assertTrue(algo.metrics.peek(f"worker_{id_}_recreated") > 1)
+        algo.stop()
+
+    def test_on_init_and_checkpoint_loaded(self):
+        config = (
+            PPOConfig()
+            .environment("CartPole-v1")
+            .callbacks(InitAndCheckpointRestoredCallbacks)
+        )
+        algo = config.build()
+        callbacks = algo.callbacks[0]
+        self.assertTrue(callbacks._on_init_was_called)
+        self.assertTrue(not hasattr(callbacks, "_on_checkpoint_loaded_was_called"))
+        algo.train()
+        # Save algo and restore.
+        with tempfile.TemporaryDirectory() as tmpdir:
+            algo.save(checkpoint_dir=tmpdir)
+            self.assertTrue(not hasattr(callbacks, "_on_checkpoint_loaded_was_called"))
+            algo.load_checkpoint(checkpoint_dir=tmpdir)
+            self.assertTrue(callbacks._on_checkpoint_loaded_was_called)
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))