chore: import upstream snapshot with attribution

python/ray/tune/tests/_test_trial_runner_pg.py

@@ -0,0 +1,318 @@
+import os
+import sys
+import time
+import unittest
+
+import numpy as np
+
+import ray
+from ray import tune
+from ray.cluster_utils import Cluster
+from ray.rllib import _register_all
+from ray.tune import Callback
+from ray.tune.execution.placement_groups import PlacementGroupFactory
+from ray.tune.execution.ray_trial_executor import RayTrialExecutor
+from ray.tune.execution.trial_runner import TrialRunner
+from ray.tune.experiment import Trial
+from ray.util import placement_group_table
+
+
+class TrialRunnerPlacementGroupTest(unittest.TestCase):
+    def setUp(self):
+        os.environ["TUNE_GLOBAL_CHECKPOINT_S"] = "10000"
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "auto"  # Reset default
+        self.head_cpus = 8
+        self.head_gpus = 4
+        self.head_custom = 16
+
+        self.cluster = Cluster(
+            initialize_head=True,
+            connect=True,
+            head_node_args={
+                "include_dashboard": False,
+                "num_cpus": self.head_cpus,
+                "num_gpus": self.head_gpus,
+                "resources": {"custom": self.head_custom},
+                "_system_config": {
+                    "health_check_initial_delay_ms": 0,
+                    "health_check_period_ms": 1000,
+                    "health_check_failure_threshold": 10,
+                },
+            },
+        )
+        # Pytest doesn't play nicely with imports
+        _register_all()
+
+    def tearDown(self):
+        ray.shutdown()
+        self.cluster.shutdown()
+        _register_all()  # re-register the evicted objects
+
+    def _assertCleanup(self, trial_executor):
+        # Assert proper cleanup
+        resource_manager = trial_executor._resource_manager
+        self.assertFalse(resource_manager._pg_to_request)
+        self.assertFalse(resource_manager._acquired_pgs)
+        self.assertFalse(resource_manager._staging_future_to_pg)
+        self.assertFalse(resource_manager._pg_to_staging_future)
+        for rr in resource_manager._request_to_staged_pgs:
+            self.assertFalse(resource_manager._request_to_staged_pgs[rr])
+        for rr in resource_manager._request_to_ready_pgs:
+            self.assertFalse(resource_manager._request_to_ready_pgs[rr])
+
+        num_non_removed_pgs = len(
+            [p for pid, p in placement_group_table().items() if p["state"] != "REMOVED"]
+        )
+        self.assertEqual(num_non_removed_pgs, 0)
+
+    def testPlacementGroupRequests(self, reuse_actors=False, scheduled=10):
+        """In this test we try to start 10 trials but only have resources
+        for 2. Placement groups should still be created and PENDING.
+
+        Eventually they should be scheduled sequentially (i.e. in pairs
+        of two)."""
+        # Since we check per-step placement groups, set the reconcilation
+        # interval to 0
+        os.environ["TUNE_PLACEMENT_GROUP_RECON_INTERVAL"] = "0"
+
+        def train_fn(config):
+            time.sleep(1)
+            now = time.time()
+            tune.report(end=now - config["start_time"])
+
+        head_bundle = {"CPU": 4, "GPU": 0, "custom": 0}
+        child_bundle = {"custom": 1}
+        # Manually calculated number of parallel trials
+        max_num_parallel = 2
+
+        placement_group_factory = PlacementGroupFactory(
+            [head_bundle, child_bundle, child_bundle]
+        )
+
+        trial_executor = RayTrialExecutor(reuse_actors=reuse_actors)
+        trial_executor.setup(max_pending_trials=max_num_parallel)
+
+        this = self
+
+        class _TestCallback(Callback):
+            def on_step_end(self, iteration, trials, **info):
+                num_finished = len(
+                    [
+                        t
+                        for t in trials
+                        if t.status == Trial.TERMINATED or t.status == Trial.ERROR
+                    ]
+                )
+
+                resource_manager = trial_executor._resource_manager
+
+                num_staging = sum(
+                    len(s) for s in resource_manager._request_to_staged_pgs.values()
+                )
+                num_ready = sum(
+                    len(s) for s in resource_manager._request_to_ready_pgs.values()
+                )
+                num_in_use = len(resource_manager._acquired_pgs)
+                num_cached = trial_executor._actor_cache.num_cached_objects
+
+                total_num_tracked = num_staging + num_ready + num_in_use + num_cached
+
+                # All trials should be scheduled
+                this.assertEqual(
+                    scheduled,
+                    min(scheduled, len(trials)),
+                    msg=f"Num trials iter {iteration}",
+                )
+
+                # The following two tests were relaxed for reuse_actors=True
+                # so that up to `max_num_parallel` more placement groups can
+                # exist than we would expect. This is because caching
+                # relies on reconciliation for cleanup to avoid overscheduling
+                # of new placement groups.
+                num_parallel_reuse = int(reuse_actors) * max_num_parallel
+
+                # The number of PGs should decrease when trials finish
+                # We allow a constant excess of 1 here because the trial will
+                # be TERMINATED and the resources only returned after the trainable
+                # cleanup future succeeded. Because num_finished will increase,
+                # this still asserts that the number of PGs goes down over time.
+                this.assertGreaterEqual(
+                    max(scheduled, len(trials)) - num_finished + 1 + num_parallel_reuse,
+                    total_num_tracked,
+                    msg=f"Num tracked iter {iteration}, {len(trials)}, "
+                    f"{scheduled}, {num_finished}, {num_parallel_reuse}",
+                )
+
+        start = time.time()
+        out = tune.run(
+            train_fn,
+            config={"start_time": start},
+            resources_per_trial=placement_group_factory,
+            num_samples=10,
+            trial_executor=trial_executor,
+            callbacks=[_TestCallback()],
+            reuse_actors=reuse_actors,
+            verbose=2,
+        )
+
+        trial_end_times = sorted(t.last_result["end"] for t in out.trials)
+        print("Trial end times:", trial_end_times)
+        max_diff = trial_end_times[-1] - trial_end_times[0]
+
+        # Not all trials have been run in parallel
+        self.assertGreater(max_diff, 3)
+
+        # Some trials should have run in parallel
+        # Todo: Re-enable when using buildkite
+        # self.assertLess(max_diff, 10)
+
+        self._assertCleanup(trial_executor)
+
+    def testPlacementGroupRequestsWithActorReuse(self):
+        """Assert that reuse actors doesn't leak placement groups"""
+        self.testPlacementGroupRequests(reuse_actors=True)
+
+    def testPlacementGroupLimitedRequests(self):
+        """Assert that maximum number of placement groups is enforced."""
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "6"
+        self.testPlacementGroupRequests(scheduled=6)
+
+    def testPlacementGroupLimitedRequestsWithActorReuse(self):
+        os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "6"
+        self.testPlacementGroupRequests(reuse_actors=True, scheduled=6)
+
+    def testPlacementGroupDistributedTraining(self, reuse_actors=False):
+        """Run distributed training using placement groups.
+
+        Each trial requests 4 CPUs and starts 4 remote training workers.
+        """
+
+        head_bundle = {"CPU": 1, "GPU": 0, "custom": 0}
+        child_bundle = {"CPU": 1}
+
+        placement_group_factory = PlacementGroupFactory(
+            [head_bundle, child_bundle, child_bundle, child_bundle]
+        )
+
+        @ray.remote
+        class TrainingActor:
+            def train(self, val):
+                time.sleep(1)
+                return val
+
+        def train_fn(config):
+            base = config["base"]
+            actors = [TrainingActor.remote() for _ in range(4)]
+            futures = [
+                actor.train.remote(base + 2 * i) for i, actor in enumerate(actors)
+            ]
+            results = ray.get(futures)
+
+            end = time.time() - config["start_time"]
+            tune.report(avg=np.mean(results), end=end)
+
+        trial_executor = RayTrialExecutor(reuse_actors=reuse_actors)
+
+        start = time.time()
+        out = tune.run(
+            train_fn,
+            config={
+                "start_time": start,
+                "base": tune.grid_search(list(range(0, 100, 10))),
+            },
+            resources_per_trial=placement_group_factory,
+            num_samples=1,
+            trial_executor=trial_executor,
+            reuse_actors=reuse_actors,
+            verbose=2,
+        )
+
+        avgs = sorted(t.last_result["avg"] for t in out.trials)
+        self.assertSequenceEqual(avgs, list(range(3, 103, 10)))
+
+        trial_end_times = sorted(t.last_result["end"] for t in out.trials)
+        print("Trial end times:", trial_end_times)
+        max_diff = trial_end_times[-1] - trial_end_times[0]
+
+        # Not all trials have been run in parallel
+        self.assertGreater(max_diff, 3)
+
+        # Some trials should have run in parallel
+        # Todo: Re-enable when using buildkite
+        # self.assertLess(max_diff, 10)
+
+        self._assertCleanup(trial_executor)
+
+    def testPlacementGroupDistributedTrainingWithActorReuse(self):
+        self.testPlacementGroupDistributedTraining(reuse_actors=True)
+
+
+class TrialRunnerPlacementGroupHeterogeneousTest(unittest.TestCase):
+    def tearDown(self) -> None:
+        if ray.is_initialized:
+            ray.shutdown()
+
+    def testResourceDeadlock(self):
+        """Tests that resource deadlock is avoided for heterogeneous PGFs.
+
+        We start 4 trials in a cluster with 2 CPUs. The first two trials
+        require 1 CPU each, the third trial 2 CPUs, the fourth trial 1 CPU.
+
+        The second trial needs a bit more time to finish. This means that the
+        resources from the first trial will be freed, and the PG of the
+        _fourth_ trial becomes ready (not that of the third trial, because that
+        requires 2 CPUs - however, one is still occupied by trial 2).
+
+        After the first two trials finished, the FIFOScheduler tries to start
+        the third trial. However, it can't be started because its placement
+        group is not ready. Instead, the placement group of the fourth
+        trial is ready. Thus, we opt to run the fourth trial instead.
+        """
+
+        def train_fn(config):
+            time.sleep(config["sleep"])
+            return 4
+
+        ray.init(num_cpus=2)
+
+        tune.register_trainable("het", train_fn)
+        pgf1 = PlacementGroupFactory([{"CPU": 1}])
+        pgf2 = PlacementGroupFactory([{"CPU": 2}])
+
+        trial1 = Trial("het", config={"sleep": 0}, placement_group_factory=pgf1)
+        trial2 = Trial("het", config={"sleep": 2}, placement_group_factory=pgf1)
+        trial3 = Trial("het", config={"sleep": 0}, placement_group_factory=pgf2)
+        trial4 = Trial("het", config={"sleep": 0}, placement_group_factory=pgf1)
+
+        runner = TrialRunner(fail_fast=True)
+        runner.add_trial(trial1)
+        runner.add_trial(trial2)
+        runner.add_trial(trial3)
+        runner.add_trial(trial4)
+
+        timeout = time.monotonic() + 30
+        while not runner.is_finished():
+            # We enforce a timeout here
+            self.assertLess(
+                time.monotonic(), timeout, msg="Ran into a resource deadlock"
+            )
+
+            runner.step()
+
+
+def test_placement_group_no_cpu_trainer():
+    """Bundles with only GPU:1 but no CPU should work"""
+    ray.init(num_gpus=1, num_cpus=1)
+    pgf = PlacementGroupFactory([{"GPU": 1, "CPU": 0}, {"CPU": 1}])
+
+    def train_fn(config):
+        time.sleep(1)
+        return 5
+
+    tune.run(train_fn, resources_per_trial=pgf)
+
+
+if __name__ == "__main__":
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))