[tune] support resume for search algorithms (#9972)

python/ray/tune/tests/_test_cluster_interrupt_searcher.py

@@ -0,0 +1,56 @@
+import argparse
+
+from ray.tune import run
+from ray.tune.examples.async_hyperband_example import MyTrainableClass
+from ray.tune.suggest.hyperopt import HyperOptSearch
+from ray.tune.suggest.suggestion import ConcurrencyLimiter
+
+from hyperopt import hp
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="PyTorch Example (FOR TEST ONLY)")
+    parser.add_argument(
+        "--resume", action="store_true", help="Resuming from checkpoint.")
+    parser.add_argument("--local-dir", help="Checkpoint path")
+    parser.add_argument(
+        "--ray-address",
+        help="Address of Ray cluster for seamless distributed execution.")
+    args = parser.parse_args()
+
+    space = {
+        "width": hp.uniform("width", 0, 20),
+        "height": hp.uniform("height", -100, 100),
+        "activation": hp.choice("activation", ["relu", "tanh"])
+    }
+    current_best_params = [
+        {
+            "width": 1,
+            "height": 2,
+            "activation": 0  # Activation will be relu
+        },
+        {
+            "width": 4,
+            "height": 2,
+            "activation": 1  # Activation will be tanh
+        }
+    ]
+    algo = HyperOptSearch(
+        space,
+        metric="episode_reward_mean",
+        mode="max",
+        random_state_seed=5,
+        points_to_evaluate=current_best_params)
+    algo = ConcurrencyLimiter(algo, max_concurrent=1)
+    from ray.tune import register_trainable
+    register_trainable("trainable", MyTrainableClass)
+    run("trainable",
+        search_alg=algo,
+        global_checkpoint_period=0,
+        resume=args.resume,
+        verbose=0,
+        num_samples=20,
+        fail_fast=True,
+        stop={"training_iteration": 2},
+        local_dir=args.local_dir,
+        name="experiment")

python/ray/tune/tests/test_api.py

@@ -23,7 +23,7 @@ from ray.tune.logger import Logger
 from ray.tune.experiment import Experiment
 from ray.tune.resources import Resources
 from ray.tune.suggest import grid_search
-from ray.tune.suggest.suggestion import _MockSuggestionAlgorithm
+from ray.tune.suggest._mock import _MockSuggestionAlgorithm
 from ray.tune.utils import (flatten_dict, get_pinned_object,
                             pin_in_object_store)
 from ray.tune.utils.mock import mock_storage_client, MOCK_REMOTE_DIR

python/ray/tune/tests/test_cluster.py

@@ -4,6 +4,7 @@ import time
 import os
 import pytest
 import shutil
+import subprocess
 import sys
 from unittest.mock import MagicMock, patch
 
@@ -713,6 +714,80 @@ tune.run(
     cluster.shutdown()
 
 
+def test_cluster_interrupt_searcher(start_connected_cluster, tmpdir):
+    """Tests restoration of HyperOptSearch experiment on cluster shutdown
+    with actual interrupt.
+
+    Restoration should restore both state of trials
+    and previous search algorithm (HyperOptSearch) state.
+    This is an end-to-end test.
+    """
+    cluster = start_connected_cluster
+    dirpath = str(tmpdir)
+    local_checkpoint_dir = os.path.join(dirpath, "experiment")
+    from ray.tune.examples.async_hyperband_example import MyTrainableClass
+    from ray.tune import register_trainable
+    register_trainable("trainable", MyTrainableClass)
+
+    def execute_script_with_args(*args):
+        current_dir = os.path.dirname(__file__)
+        script = os.path.join(current_dir,
+                              "_test_cluster_interrupt_searcher.py")
+        subprocess.Popen([sys.executable, script] + list(args))
+
+    args = ["--ray-address", cluster.address, "--local-dir", dirpath]
+    execute_script_with_args(*args)
+    # Wait until the right checkpoint is saved.
+    # The trainable returns every 0.5 seconds, so this should not miss
+    # the checkpoint.
+    for i in range(50):
+        if TrialRunner.checkpoint_exists(local_checkpoint_dir):
+            # Inspect the internal trialrunner
+            runner = TrialRunner(
+                resume="LOCAL", local_checkpoint_dir=local_checkpoint_dir)
+            trials = runner.get_trials()
+            if trials and len(trials) >= 10:
+                break
+        time.sleep(.5)
+
+    if not TrialRunner.checkpoint_exists(local_checkpoint_dir):
+        raise RuntimeError(
+            f"Checkpoint file didn't appear in {local_checkpoint_dir}. "
+            f"Current list: {os.listdir(local_checkpoint_dir)}.")
+
+    ray.shutdown()
+    cluster.shutdown()
+
+    cluster = _start_new_cluster()
+    execute_script_with_args(*(args + ["--resume"]))
+
+    time.sleep(2)
+
+    register_trainable("trainable", MyTrainableClass)
+    reached = False
+    for i in range(50):
+        if TrialRunner.checkpoint_exists(local_checkpoint_dir):
+            # Inspect the internal trialrunner
+            runner = TrialRunner(
+                resume="LOCAL", local_checkpoint_dir=local_checkpoint_dir)
+            trials = runner.get_trials()
+            if len(trials) == 0:
+                continue  # nonblocking script hasn't resumed yet, wait
+            reached = True
+            assert len(trials) >= 10
+            assert len(trials) <= 20
+            if len(trials) == 20:
+                break
+            else:
+                stop_fn = runner.trial_executor.stop_trial
+                [stop_fn(t) for t in trials if t.status is not Trial.ERROR]
+        time.sleep(.5)
+    assert reached is True
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
 if __name__ == "__main__":
     import pytest
     sys.exit(pytest.main(["-v", __file__]))

python/ray/tune/tests/test_trial_runner_3.py

@@ -1,4 +1,6 @@
+from collections import Counter
 import os
+import pickle
 import shutil
 import sys
 import tempfile
@@ -14,14 +16,17 @@ from ray.tune.trial import Trial
 from ray.tune.trial_runner import TrialRunner
 from ray.tune.resources import Resources, json_to_resources, resources_to_json
 from ray.tune.suggest.repeater import Repeater
-from ray.tune.suggest.suggestion import (_MockSuggestionAlgorithm,
-                                         SearchGenerator, Searcher)
+from ray.tune.suggest._mock import _MockSuggestionAlgorithm
+from ray.tune.suggest.suggestion import Searcher, ConcurrencyLimiter
+from ray.tune.suggest.search_generator import SearchGenerator
 
 
 class TrialRunnerTest3(unittest.TestCase):
     def tearDown(self):
         ray.shutdown()
         _register_all()  # re-register the evicted objects
+        if "CUDA_VISIBLE_DEVICES" in os.environ:
+            del os.environ["CUDA_VISIBLE_DEVICES"]
 
     def testStepHook(self):
         ray.init(num_cpus=4, num_gpus=2)
@@ -264,6 +269,92 @@ class TrialRunnerTest3(unittest.TestCase):
         self.assertTrue(searcher.is_finished())
         self.assertRaises(TuneError, runner.step)
 
+    def testSearcherSaveRestore(self):
+        ray.init(num_cpus=8, local_mode=True)
+        tmpdir = tempfile.mkdtemp()
+
+        def create_searcher():
+            class TestSuggestion(Searcher):
+                def __init__(self, index):
+                    self.index = index
+                    self.returned_result = []
+                    super().__init__(metric="result", mode="max")
+
+                def suggest(self, trial_id):
+                    self.index += 1
+                    return {"test_variable": self.index}
+
+                def on_trial_complete(self, trial_id, result=None, **kwargs):
+                    self.returned_result.append(result)
+
+                def save(self, checkpoint_path):
+                    with open(checkpoint_path, "wb") as f:
+                        pickle.dump(self.__dict__, f)
+
+                def restore(self, checkpoint_path):
+                    with open(checkpoint_path, "rb") as f:
+                        self.__dict__.update(pickle.load(f))
+
+            searcher = TestSuggestion(0)
+            searcher = ConcurrencyLimiter(searcher, max_concurrent=2)
+            searcher = Repeater(searcher, repeat=3, set_index=False)
+            search_alg = SearchGenerator(searcher)
+            experiment_spec = {
+                "run": "__fake",
+                "num_samples": 20,
+                "stop": {
+                    "training_iteration": 2
+                }
+            }
+            experiments = [Experiment.from_json("test", experiment_spec)]
+            search_alg.add_configurations(experiments)
+            return search_alg
+
+        searcher = create_searcher()
+        runner = TrialRunner(
+            search_alg=searcher,
+            local_checkpoint_dir=tmpdir,
+            checkpoint_period=-1)
+        for i in range(6):
+            runner.step()
+
+        assert len(
+            runner.get_trials()) == 6, [t.config for t in runner.get_trials()]
+        runner.checkpoint()
+        trials = runner.get_trials()
+        [
+            runner.trial_executor.stop_trial(t) for t in trials
+            if t.status is not Trial.ERROR
+        ]
+        del runner
+        # stop_all(runner.get_trials())
+
+        searcher = create_searcher()
+        runner2 = TrialRunner(
+            search_alg=searcher, local_checkpoint_dir=tmpdir, resume="LOCAL")
+        assert len(runner2.get_trials()) == 6, [
+            t.config for t in runner2.get_trials()
+        ]
+
+        def trial_statuses():
+            return [t.status for t in runner2.get_trials()]
+
+        def num_running_trials():
+            return sum(t.status == Trial.RUNNING for t in runner2.get_trials())
+
+        for i in range(6):
+            runner2.step()
+        assert len(set(trial_statuses())) == 1
+        assert Trial.RUNNING in trial_statuses()
+        for i in range(20):
+            runner2.step()
+            assert 1 <= num_running_trials() <= 6
+        evaluated = [
+            t.evaluated_params["test_variable"] for t in runner2.get_trials()
+        ]
+        count = Counter(evaluated)
+        assert all(v <= 3 for v in count.values())
+
     def testTrialSaveRestore(self):
         """Creates different trials to test runner.checkpoint/restore."""
         ray.init(num_cpus=3)
@@ -512,6 +603,90 @@ class SearchAlgorithmTest(unittest.TestCase):
         parameter_set = {t.evaluated_params["test_variable"] for t in trials}
         self.assertEquals(len(parameter_set), 3)
 
+    def testSetGetRepeater(self):
+        ray.init(num_cpus=4)
+
+        class TestSuggestion(Searcher):
+            def __init__(self, index):
+                self.index = index
+                self.returned_result = []
+                super().__init__(metric="result", mode="max")
+
+            def suggest(self, trial_id):
+                self.index += 1
+                return {"score": self.index}
+
+            def on_trial_complete(self, trial_id, result=None, **kwargs):
+                self.returned_result.append(result)
+
+        searcher = TestSuggestion(0)
+        repeater1 = Repeater(searcher, repeat=3, set_index=False)
+        for i in range(3):
+            assert repeater1.suggest(f"test_{i}")["score"] == 1
+        for i in range(2):  # An incomplete set of results
+            assert repeater1.suggest(f"test_{i}_2")["score"] == 2
+
+        # Restore a new one
+        state = repeater1.get_state()
+        del repeater1
+        new_repeater = Repeater(searcher, repeat=1, set_index=True)
+        new_repeater.set_state(state)
+        assert new_repeater.repeat == 3
+        assert new_repeater.suggest("test_2_2")["score"] == 2
+        assert new_repeater.suggest("test_x")["score"] == 3
+
+        # Report results
+        for i in range(3):
+            new_repeater.on_trial_complete(f"test_{i}", {"result": 2})
+
+        for i in range(3):
+            new_repeater.on_trial_complete(f"test_{i}_2", {"result": -i * 10})
+
+        assert len(new_repeater.searcher.returned_result) == 2
+        assert new_repeater.searcher.returned_result[-1] == {"result": -10}
+
+        # Finish the rest of the last trial group
+        new_repeater.on_trial_complete("test_x", {"result": 3})
+        assert new_repeater.suggest("test_y")["score"] == 3
+        new_repeater.on_trial_complete("test_y", {"result": 3})
+        assert len(new_repeater.searcher.returned_result) == 2
+        assert new_repeater.suggest("test_z")["score"] == 3
+        new_repeater.on_trial_complete("test_z", {"result": 3})
+        assert len(new_repeater.searcher.returned_result) == 3
+        assert new_repeater.searcher.returned_result[-1] == {"result": 3}
+
+    def testSetGetLimiter(self):
+        ray.init(num_cpus=4)
+
+        class TestSuggestion(Searcher):
+            def __init__(self, index):
+                self.index = index
+                self.returned_result = []
+                super().__init__(metric="result", mode="max")
+
+            def suggest(self, trial_id):
+                self.index += 1
+                return {"score": self.index}
+
+            def on_trial_complete(self, trial_id, result=None, **kwargs):
+                self.returned_result.append(result)
+
+        searcher = TestSuggestion(0)
+        limiter = ConcurrencyLimiter(searcher, max_concurrent=2)
+        assert limiter.suggest("test_1")["score"] == 1
+        assert limiter.suggest("test_2")["score"] == 2
+        assert limiter.suggest("test_3") is None
+
+        state = limiter.get_state()
+        del limiter
+        limiter2 = ConcurrencyLimiter(searcher, max_concurrent=3)
+        limiter2.set_state(state)
+        assert limiter2.suggest("test_4") is None
+        assert limiter2.suggest("test_5") is None
+        limiter2.on_trial_complete("test_1", {"result": 3})
+        limiter2.on_trial_complete("test_2", {"result": 3})
+        assert limiter2.suggest("test_3")["score"] == 3
+
 
 class ResourcesTest(unittest.TestCase):
     def testSubtraction(self):