chore: import upstream snapshot with attribution

rllib/algorithms/tests/test_placement_groups.py

@@ -0,0 +1,133 @@
+import os
+import unittest
+
+import ray
+from ray import tune
+from ray.rllib.algorithms.ppo import PPO, PPOConfig
+from ray.tune import Callback
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.experiment import Trial
+from ray.tune.result import TRAINING_ITERATION
+
+trial_executor = None
+
+
+class _TestCallback(Callback):
+    def on_step_end(self, iteration, trials, **info):
+        num_running = len([t for t in trials if t.status == Trial.RUNNING])
+
+        # All 3 trials (3 different learning rates) should be scheduled.
+        assert 3 == min(3, len(trials))
+        # Cannot run more than 2 at a time
+        # (due to different resource restrictions in the test cases).
+        assert num_running <= 2
+
+
+class TestPlacementGroups(unittest.TestCase):
+    def setUp(self) -> None:
+        os.environ["TUNE_PLACEMENT_GROUP_RECON_INTERVAL"] = "0"
+        ray.init(num_cpus=6)
+
+    def tearDown(self) -> None:
+        ray.shutdown()
+
+    def test_overriding_default_resource_request(self):
+        # 3 Trials: Can only run 2 at a time (num_cpus=6; needed: 3).
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_env_runner_and_connector_v2=False,
+                enable_rl_module_and_learner=False,
+            )
+            .training(
+                model={"fcnet_hiddens": [10]}, lr=tune.grid_search([0.1, 0.01, 0.001])
+            )
+            .environment("CartPole-v1")
+            .env_runners(num_env_runners=2)
+            .framework("tf")
+        )
+
+        # Create an Algorithm with an overridden default_resource_request
+        # method that returns a PlacementGroupFactory.
+
+        class MyAlgo(PPO):
+            @classmethod
+            def default_resource_request(cls, config):
+                head_bundle = {"CPU": 1, "GPU": 0}
+                child_bundle = {"CPU": 1}
+                return PlacementGroupFactory(
+                    [head_bundle, child_bundle, child_bundle],
+                    strategy=config["placement_strategy"],
+                )
+
+        tune.register_trainable("my_trainable", MyAlgo)
+
+        tune.Tuner(
+            "my_trainable",
+            param_space=config,
+            run_config=tune.RunConfig(
+                stop={TRAINING_ITERATION: 2},
+                verbose=2,
+                callbacks=[_TestCallback()],
+            ),
+        ).fit()
+
+    def test_default_resource_request(self):
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_env_runner_and_connector_v2=False,
+                enable_rl_module_and_learner=False,
+            )
+            .resources(placement_strategy="SPREAD")
+            .env_runners(
+                num_env_runners=2,
+                num_cpus_per_env_runner=2,
+            )
+            .training(
+                model={"fcnet_hiddens": [10]}, lr=tune.grid_search([0.1, 0.01, 0.001])
+            )
+            .environment("CartPole-v1")
+            .framework("torch")
+        )
+        # 3 Trials: Can only run 1 at a time (num_cpus=6; needed: 5).
+
+        tune.Tuner(
+            PPO,
+            param_space=config,
+            run_config=tune.RunConfig(
+                stop={TRAINING_ITERATION: 2},
+                verbose=2,
+                callbacks=[_TestCallback()],
+            ),
+            tune_config=tune.TuneConfig(reuse_actors=False),
+        ).fit()
+
+    def test_default_resource_request_plus_manual_leads_to_error(self):
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_env_runner_and_connector_v2=False,
+                enable_rl_module_and_learner=False,
+            )
+            .training(model={"fcnet_hiddens": [10]})
+            .environment("CartPole-v1")
+            .env_runners(num_env_runners=0)
+        )
+
+        try:
+            tune.Tuner(
+                tune.with_resources(PPO, PlacementGroupFactory([{"CPU": 1}])),
+                param_space=config,
+                run_config=tune.RunConfig(stop={TRAINING_ITERATION: 2}, verbose=2),
+            ).fit()
+        except ValueError as e:
+            assert "have been automatically set to" in e.args[0]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))