[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):