[tune] Save/Restore for Suggestion Algs (#5719)

python/ray/tune/suggest/hyperopt.py

@@ -209,11 +209,11 @@ class HyperOptSearch(SuggestionAlgorithm):
 
     def save(self, checkpoint_dir):
         trials_object = (self._hpopt_trials, self.rstate.get_state())
-        with open(checkpoint_dir, "wb") as output:
-            pickle.dump(trials_object, output)
+        with open(checkpoint_dir, "wb") as outputFile:
+            pickle.dump(trials_object, outputFile)
 
     def restore(self, checkpoint_dir):
-        with open(checkpoint_dir, "rb") as input:
-            trials_object = pickle.load(input)
+        with open(checkpoint_dir, "rb") as inputFile:
+            trials_object = pickle.load(inputFile)
         self._hpopt_trials = trials_object[0]
         self.rstate.set_state(trials_object[1])

python/ray/tune/suggest/nevergrad.py

@@ -3,6 +3,7 @@ from __future__ import division
 from __future__ import print_function
 
 import logging
+import pickle
 try:
     import nevergrad as ng
 except ImportError:
@@ -145,3 +146,14 @@ class NevergradSearch(SuggestionAlgorithm):
 
     def _num_live_trials(self):
         return len(self._live_trial_mapping)
+
+    def save(self, checkpoint_dir):
+        trials_object = (self._nevergrad_opt, self._parameters)
+        with open(checkpoint_dir, "wb") as outputFile:
+            pickle.dump(trials_object, outputFile)
+
+    def restore(self, checkpoint_dir):
+        with open(checkpoint_dir, "rb") as inputFile:
+            trials_object = pickle.load(inputFile)
+        self._nevergrad_opt = trials_object[0]
+        self._parameters = trials_object[1]

python/ray/tune/suggest/sigopt.py

@@ -5,6 +5,7 @@ from __future__ import print_function
 import copy
 import os
 import logging
+import pickle
 try:
     import sigopt as sgo
 except ImportError:
@@ -140,3 +141,14 @@ class SigOptSearch(SuggestionAlgorithm):
 
     def _num_live_trials(self):
         return len(self._live_trial_mapping)
+
+    def save(self, checkpoint_dir):
+        trials_object = (self.conn, self.experiment)
+        with open(checkpoint_dir, "wb") as outputFile:
+            pickle.dump(trials_object, outputFile)
+
+    def restore(self, checkpoint_dir):
+        with open(checkpoint_dir, "rb") as inputFile:
+            trials_object = pickle.load(inputFile)
+        self.conn = trials_object[0]
+        self.experiment = trials_object[1]

python/ray/tune/suggest/skopt.py

@@ -160,11 +160,12 @@ class SkOptSearch(SuggestionAlgorithm):
         return len(self._live_trial_mapping)
 
     def save(self, checkpoint_dir):
-        trials_object = self._skopt_opt
-        with open(checkpoint_dir, "wb") as output:
-            pickle.dump(trials_object, output)
+        trials_object = (self._initial_points, self._skopt_opt)
+        with open(checkpoint_dir, "wb") as outputFile:
+            pickle.dump(trials_object, outputFile)
 
     def restore(self, checkpoint_dir):
-        with open(checkpoint_dir, "rb") as input:
-            trials_object = pickle.load(input)
-        self._skopt_opt = trials_object
+        with open(checkpoint_dir, "rb") as inputFile:
+            trials_object = pickle.load(inputFile)
+        self._initial_points = trials_object[0]
+        self._skopt_opt = trials_object[1]

python/ray/tune/suggest/suggestion.py

@@ -129,6 +129,12 @@ class SuggestionAlgorithm(SearchAlgorithm):
         """
         raise NotImplementedError
 
+    def save(self, checkpoint_dir):
+        raise NotImplementedError
+
+    def restore(self, checkpoint_dir):
+        raise NotImplementedError
+
 
 class _MockSuggestionAlgorithm(SuggestionAlgorithm):
     def __init__(self, max_concurrent=2, **kwargs):

python/ray/tune/tests/test_tune_restore.py

@@ -8,6 +8,8 @@ import os
 import shutil
 import tempfile
 import unittest
+import skopt
+import numpy as np
 
 import ray
 from ray import tune
@@ -16,6 +18,10 @@ from ray.tune.util import validate_save_restore
 from ray.rllib import _register_all
 from ray.tune.suggest.hyperopt import HyperOptSearch
 from ray.tune.suggest.bayesopt import BayesOptSearch
+from ray.tune.suggest.skopt import SkOptSearch
+from ray.tune.suggest.nevergrad import NevergradSearch
+from nevergrad.optimization import optimizerlib
+from ray.tune.suggest.sigopt import SigOptSearch
 
 
 class TuneRestoreTest(unittest.TestCase):
@@ -115,7 +121,7 @@ class AutoInitTest(unittest.TestCase):
         _register_all()
 
 
-class HyperoptWarmStartTest(unittest.TestCase):
+class AbstractWarmStartTest(object):
     def setUp(self):
         ray.init(local_mode=True)
         self.tmpdir = tempfile.mkdtemp()
@@ -125,6 +131,38 @@ class HyperoptWarmStartTest(unittest.TestCase):
         ray.shutdown()
         _register_all()
 
+    def set_basic_conf(self):
+        raise NotImplementedError()
+
+    def run_exp_1(self):
+        np.random.seed(162)
+        search_alg, cost = self.set_basic_conf()
+        results_exp_1 = tune.run(cost, num_samples=15, search_alg=search_alg)
+        self.log_dir = os.path.join(self.tmpdir, "warmStartTest.pkl")
+        search_alg.save(self.log_dir)
+        return results_exp_1
+
+    def run_exp_2(self):
+        search_alg2, cost = self.set_basic_conf()
+        search_alg2.restore(self.log_dir)
+        return tune.run(cost, num_samples=15, search_alg=search_alg2)
+
+    def run_exp_3(self):
+        np.random.seed(162)
+        search_alg3, cost = self.set_basic_conf()
+        return tune.run(cost, num_samples=30, search_alg=search_alg3)
+
+    def testWarmStart(self):
+        results_exp_1 = self.run_exp_1()
+        results_exp_2 = self.run_exp_2()
+        results_exp_3 = self.run_exp_3()
+        trials_1_config = [trial.config for trial in results_exp_1.trials]
+        trials_2_config = [trial.config for trial in results_exp_2.trials]
+        trials_3_config = [trial.config for trial in results_exp_3.trials]
+        self.assertEqual(trials_1_config + trials_2_config, trials_3_config)
+
+
+class HyperoptWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
     def set_basic_conf(self):
         space = {
             "x": hp.uniform("x", 0, 10),
@@ -144,48 +182,13 @@ class HyperoptWarmStartTest(unittest.TestCase):
             random_state_seed=5)
         return search_alg, cost
 
-    def run_exp_1(self):
-        search_alg, cost = self.set_basic_conf()
-        results_exp_1 = tune.run(cost, num_samples=15, search_alg=search_alg)
-        self.log_dir = os.path.join(self.tmpdir, "trials_algo_hyo.pkl")
-        search_alg.save(self.log_dir)
-        return results_exp_1
-
-    def run_exp_2(self):
-        search_alg2, cost = self.set_basic_conf()
-        search_alg2.restore(self.log_dir)
-        return tune.run(cost, num_samples=15, search_alg=search_alg2)
-
-    def run_exp_3(self):
-        search_alg3, cost = self.set_basic_conf()
-        return tune.run(cost, num_samples=30, search_alg=search_alg3)
-
-    def testHyperoptWarmStart(self):
-        results_exp_1 = self.run_exp_1()
-        results_exp_2 = self.run_exp_2()
-        results_exp_3 = self.run_exp_3()
-        trials_1_config = [trial.config for trial in results_exp_1.trials]
-        trials_2_config = [trial.config for trial in results_exp_2.trials]
-        trials_3_config = [trial.config for trial in results_exp_3.trials]
-        self.assertEqual(trials_1_config + trials_2_config, trials_3_config)
-
-
-class BayesoptWarmStartTest(unittest.TestCase):
-    def setUp(self):
-        ray.init(local_mode=True)
-        self.tmpdir = tempfile.mkdtemp()
-
-    def tearDown(self):
-        shutil.rmtree(self.tmpdir)
-        ray.shutdown()
-        _register_all()
 
+class BayesoptWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
     def set_basic_conf(self):
         space = {"width": (0, 20), "height": (-100, 100)}
 
         def cost(space, reporter):
-            loss = space["width"]**2 + space["height"]**2
-            reporter(loss=loss)
+            reporter(loss=(space["height"] - 14)**2 - abs(space["width"] - 3))
 
         search_alg = BayesOptSearch(
             space,
@@ -199,30 +202,83 @@ class BayesoptWarmStartTest(unittest.TestCase):
             })
         return search_alg, cost
 
-    def run_exp_1(self):
-        search_alg, cost = self.set_basic_conf()
-        results_exp_1 = tune.run(cost, num_samples=15, search_alg=search_alg)
-        self.log_dir = os.path.join(self.tmpdir, "trials_algo_byo.pkl")
-        search_alg.save(self.log_dir)
-        return results_exp_1
 
-    def run_exp_2(self):
-        search_alg2, cost = self.set_basic_conf()
-        search_alg2.restore(self.log_dir)
-        return tune.run(cost, num_samples=15, search_alg=search_alg2)
+class SkoptWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        optimizer = skopt.Optimizer([(0, 20), (-100, 100)])
+        previously_run_params = [[10, 0], [15, -20]]
+        known_rewards = [-189, -1144]
 
-    def run_exp_3(self):
-        search_alg3, cost = self.set_basic_conf()
-        return tune.run(cost, num_samples=30, search_alg=search_alg3)
+        def cost(space, reporter):
+            reporter(loss=(space["height"]**2 + space["width"]**2))
 
-    def testBayesoptWarmStart(self):
-        results_exp_1 = self.run_exp_1()
-        results_exp_2 = self.run_exp_2()
-        results_exp_3 = self.run_exp_3()
-        trials_1_config = [trial.config for trial in results_exp_1.trials]
-        trials_2_config = [trial.config for trial in results_exp_2.trials]
-        trials_3_config = [trial.config for trial in results_exp_3.trials]
-        self.assertEqual(trials_1_config + trials_2_config, trials_3_config)
+        search_alg = SkOptSearch(
+            optimizer, ["width", "height"],
+            max_concurrent=1,
+            metric="loss",
+            mode="min",
+            points_to_evaluate=previously_run_params,
+            evaluated_rewards=known_rewards)
+        return search_alg, cost
+
+
+class NevergradWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        instrumentation = 2
+        parameter_names = ["height", "width"]
+        optimizer = optimizerlib.OnePlusOne(instrumentation)
+
+        def cost(space, reporter):
+            reporter(
+                mean_loss=(space["height"] - 14)**2 - abs(space["width"] - 3))
+
+        search_alg = NevergradSearch(
+            optimizer,
+            parameter_names,
+            max_concurrent=1,
+            metric="mean_loss",
+            mode="min")
+        return search_alg, cost
+
+
+class SigOptWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
+    def set_basic_conf(self):
+        space = [
+            {
+                "name": "width",
+                "type": "int",
+                "bounds": {
+                    "min": 0,
+                    "max": 20
+                },
+            },
+            {
+                "name": "height",
+                "type": "int",
+                "bounds": {
+                    "min": -100,
+                    "max": 100
+                },
+            },
+        ]
+
+        def cost(space, reporter):
+            reporter(
+                mean_loss=(space["height"] - 14)**2 - abs(space["width"] - 3))
+
+        search_alg = SigOptSearch(
+            space,
+            name="SigOpt Example Experiment",
+            max_concurrent=1,
+            metric="mean_loss",
+            mode="min")
+        return search_alg, cost
+
+    def testWarmStart(self):
+        if ("SIGOPT_KEY" not in os.environ):
+            return
+
+        super().testWarmStart()
 
 
 if __name__ == "__main__":