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

python/ray/tune/suggest/__init__.py

@@ -1,7 +1,7 @@
 from ray.tune.suggest.search import SearchAlgorithm
 from ray.tune.suggest.basic_variant import BasicVariantGenerator
-from ray.tune.suggest.suggestion import (SearchGenerator, Searcher,
-                                         ConcurrencyLimiter)
+from ray.tune.suggest.suggestion import Searcher, ConcurrencyLimiter
+from ray.tune.suggest.search_generator import SearchGenerator
 from ray.tune.suggest.variant_generator import grid_search
 from ray.tune.suggest.repeater import Repeater
 

python/ray/tune/suggest/_mock.py

@@ -0,0 +1,49 @@
+from ray.tune.suggest.suggestion import Searcher, ConcurrencyLimiter
+from ray.tune.suggest.search_generator import SearchGenerator
+
+
+class _MockSearcher(Searcher):
+    def __init__(self, **kwargs):
+        self.live_trials = {}
+        self.counter = {"result": 0, "complete": 0}
+        self.final_results = []
+        self.stall = False
+        self.results = []
+        super(_MockSearcher, self).__init__(**kwargs)
+
+    def suggest(self, trial_id):
+        if not self.stall:
+            self.live_trials[trial_id] = 1
+            return {"test_variable": 2}
+        return None
+
+    def on_trial_result(self, trial_id, result):
+        self.counter["result"] += 1
+        self.results += [result]
+
+    def on_trial_complete(self, trial_id, result=None, error=False):
+        self.counter["complete"] += 1
+        if result:
+            self._process_result(result)
+        if trial_id in self.live_trials:
+            del self.live_trials[trial_id]
+
+    def _process_result(self, result):
+        self.final_results += [result]
+
+
+class _MockSuggestionAlgorithm(SearchGenerator):
+    def __init__(self, max_concurrent=None, **kwargs):
+        self.searcher = _MockSearcher(**kwargs)
+        if max_concurrent:
+            self.searcher = ConcurrencyLimiter(
+                self.searcher, max_concurrent=max_concurrent)
+        super(_MockSuggestionAlgorithm, self).__init__(self.searcher)
+
+    @property
+    def live_trials(self):
+        return self.searcher.live_trials
+
+    @property
+    def results(self):
+        return self.searcher.results

python/ray/tune/suggest/hyperopt.py

@@ -212,13 +212,26 @@ class HyperOptSearch(Searcher):
             t for t in self._hpopt_trials.trials if t["tid"] == hyperopt_tid
         ][0]
 
+    def get_state(self):
+        return {
+            "hyperopt_trials": self._hpopt_trials,
+            "rstate": self.rstate.get_state()
+        }
+
+    def set_state(self, state):
+        self._hpopt_trials = state["hyperopt_trials"]
+        self.rstate.set_state(state["rstate"])
+
     def save(self, checkpoint_path):
-        trials_object = (self._hpopt_trials, self.rstate.get_state())
         with open(checkpoint_path, "wb") as outputFile:
-            pickle.dump(trials_object, outputFile)
+            pickle.dump(self.get_state(), outputFile)
 
     def restore(self, checkpoint_path):
         with open(checkpoint_path, "rb") as inputFile:
             trials_object = pickle.load(inputFile)
-        self._hpopt_trials = trials_object[0]
-        self.rstate.set_state(trials_object[1])
+
+        if isinstance(trials_object, tuple):
+            self._hpopt_trials = trials_object[0]
+            self.rstate.set_state(trials_object[1])
+        else:
+            self.set_state(trials_object)

python/ray/tune/suggest/repeater.py

@@ -44,7 +44,7 @@ class _TrialGroup:
 
     def add(self, trial_id):
         assert len(self._trials) < self.max_trials
-        self._trials[trial_id] = None
+        self._trials.setdefault(trial_id, None)
 
     def full(self):
         return len(self._trials) == self.max_trials
@@ -56,7 +56,8 @@ class _TrialGroup:
         self._trials[trial_id] = score
 
     def finished_reporting(self):
-        return None not in self._trials.values()
+        return None not in self._trials.values() and len(
+            self._trials) == self.max_trials
 
     def scores(self):
         return list(self._trials.values())
@@ -159,8 +160,10 @@ class Repeater(Searcher):
                 result={self.searcher.metric: np.nanmean(scores)},
                 **kwargs)
 
-    def save(self, path):
-        self.searcher.save(path)
+    def get_state(self):
+        self_state = self.__dict__.copy()
+        del self_state["searcher"]
+        return self_state
 
-    def restore(self, path):
-        self.searcher.restore(path)
+    def set_state(self, state):
+        self.__dict__.update(state)

python/ray/tune/suggest/search.py

@@ -63,8 +63,14 @@ class SearchAlgorithm:
         """Marks the search algorithm as finished."""
         self._finished = True
 
-    def save(self, *args):
+    def has_checkpoint(self, dirpath):
+        """Should return False if not restoring is not implemented."""
+        return False
+
+    def save_to_dir(self, dirpath, **kwargs):
+        """Saves a search algorithm."""
         pass
 
-    def restore(self, *args):
+    def restore_from_dir(self, dirpath):
+        """Restores a search algorithm along with its wrapped state."""
         pass

python/ray/tune/suggest/search_generator.py

@@ -0,0 +1,222 @@
+import pickle
+import os
+import copy
+import logging
+import glob
+
+from ray.tune.error import TuneError
+from ray.tune.experiment import convert_to_experiment_list
+from ray.tune.config_parser import make_parser, create_trial_from_spec
+from ray.tune.suggest.search import SearchAlgorithm
+from ray.tune.suggest.suggestion import Searcher
+from ray.tune.suggest.variant_generator import format_vars, resolve_nested_dict
+from ray.tune.trial import Trial
+from ray.tune.utils import flatten_dict, merge_dicts
+
+logger = logging.getLogger(__name__)
+
+
+def _warn_on_repeater(searcher, total_samples):
+    from ray.tune.suggest.repeater import _warn_num_samples
+    _warn_num_samples(searcher, total_samples)
+
+
+def _atomic_save(state, checkpoint_dir, file_name):
+    """Atomically saves the object to the checkpoint directory
+
+    This is automatically used by tune.run during a Tune job.
+    """
+    tmp_search_ckpt_path = os.path.join(checkpoint_dir,
+                                        ".tmp_search_generator_ckpt")
+    with open(tmp_search_ckpt_path, "wb") as f:
+        pickle.dump(state, f)
+
+    os.rename(tmp_search_ckpt_path, os.path.join(checkpoint_dir, file_name))
+
+
+def _find_newest_ckpt(dirpath, pattern):
+    """Returns path to most recently modified checkpoint."""
+    full_paths = glob.glob(os.path.join(dirpath, pattern))
+    if not full_paths:
+        return
+    most_recent_checkpoint = max(full_paths)
+    with open(most_recent_checkpoint, "rb") as f:
+        search_alg_state = pickle.load(f)
+    return search_alg_state
+
+
+class SearchGenerator(SearchAlgorithm):
+    """Generates trials to be passed to the TrialRunner.
+
+    Uses the provided ``searcher`` object to generate trials. This class
+    transparently handles repeating trials with score aggregation
+    without embedding logic into the Searcher.
+
+    Args:
+        searcher: Search object that subclasses the Searcher base class. This
+            is then used for generating new hyperparameter samples.
+    """
+    CKPT_FILE_TMPL = "search_gen_state-{}.json"
+
+    def __init__(self, searcher):
+        assert issubclass(
+            type(searcher),
+            Searcher), ("Searcher should be subclassing Searcher.")
+        self.searcher = searcher
+        self._parser = make_parser()
+        self._experiment = None
+        self._counter = 0  # Keeps track of number of trials created.
+        self._total_samples = None  # int: total samples to evaluate.
+        self._finished = False
+
+    def add_configurations(self, experiments):
+        """Registers experiment specifications.
+
+        Arguments:
+            experiments (Experiment | list | dict): Experiments to run.
+        """
+        assert not self._experiment
+        logger.debug("added configurations")
+        experiment_list = convert_to_experiment_list(experiments)
+        assert len(experiment_list) == 1, (
+            "SearchAlgorithms can only support 1 experiment at a time.")
+        self._experiment = experiment_list[0]
+        experiment_spec = self._experiment.spec
+        self._total_samples = self._experiment.spec.get("num_samples", 1)
+
+        _warn_on_repeater(self.searcher, self._total_samples)
+        if "run" not in experiment_spec:
+            raise TuneError("Must specify `run` in {}".format(experiment_spec))
+
+    def next_trials(self):
+        """Provides a batch of Trial objects to be queued into the TrialRunner.
+
+        Returns:
+            List[Trial]: A list of trials for the Runner to consume.
+        """
+        trials = []
+        while not self.is_finished():
+            trial = self.create_trial_if_possible(self._experiment.spec,
+                                                  self._experiment.name)
+            if trial is None:
+                break
+            trials.append(trial)
+        return trials
+
+    def create_trial_if_possible(self, experiment_spec, output_path):
+        logger.debug("creating trial")
+        trial_id = Trial.generate_id()
+        suggested_config = self.searcher.suggest(trial_id)
+        if suggested_config == Searcher.FINISHED:
+            self._finished = True
+            logger.debug("Searcher has finished.")
+            return
+
+        if suggested_config is None:
+            return
+        spec = copy.deepcopy(experiment_spec)
+        spec["config"] = merge_dicts(spec["config"],
+                                     copy.deepcopy(suggested_config))
+
+        # Create a new trial_id if duplicate trial is created
+        flattened_config = resolve_nested_dict(spec["config"])
+        self._counter += 1
+        tag = "{0}_{1}".format(
+            str(self._counter), format_vars(flattened_config))
+        trial = create_trial_from_spec(
+            spec,
+            output_path,
+            self._parser,
+            evaluated_params=flatten_dict(suggested_config),
+            experiment_tag=tag,
+            trial_id=trial_id)
+        return trial
+
+    def on_trial_result(self, trial_id, result):
+        """Notifies the underlying searcher."""
+        self.searcher.on_trial_result(trial_id, result)
+
+    def on_trial_complete(self, trial_id, result=None, error=False):
+        self.searcher.on_trial_complete(
+            trial_id=trial_id, result=result, error=error)
+
+    def is_finished(self):
+        return self._counter >= self._total_samples or self._finished
+
+    def get_state(self):
+        return {
+            "counter": self._counter,
+            "total_samples": self._total_samples,
+            "finished": self._finished,
+            "experiment": self._experiment
+        }
+
+    def set_state(self, state):
+        self._counter = state["counter"]
+        self._total_samples = state["total_samples"]
+        self._finished = state["finished"]
+        self._experiment = state["experiment"]
+
+    def has_checkpoint(self, dirpath):
+        return bool(
+            _find_newest_ckpt(dirpath, self.CKPT_FILE_TMPL.format("*")))
+
+    def save_to_dir(self, dirpath, session_str):
+        """Saves self + searcher to dir.
+
+        Separates the "searcher" from its wrappers (concurrency, repeating).
+        This allows the user to easily restore a given searcher.
+
+        The save operation is atomic (write/swap).
+
+        Args:
+            dirpath (str): Filepath to experiment dir.
+            session_str (str): Unique identifier of the current run
+                session.
+        """
+        searcher = self.searcher
+        search_alg_state = self.get_state()
+        while hasattr(searcher, "searcher"):
+            searcher_name = type(searcher).__name__
+            if searcher_name in search_alg_state:
+                logger.warning(
+                    "There was a duplicate when saving {}. "
+                    "Restore may not work properly.".format(searcher_name))
+            else:
+                search_alg_state["name:" +
+                                 searcher_name] = searcher.get_state()
+            searcher = searcher.searcher
+        base_searcher = searcher
+        # We save the base searcher separately for users to easily
+        # separate the searcher.
+        base_searcher.save_to_dir(dirpath, session_str)
+        _atomic_save(search_alg_state, dirpath,
+                     self.CKPT_FILE_TMPL.format(session_str))
+
+    def restore_from_dir(self, dirpath):
+        """Restores self + searcher + search wrappers from dirpath."""
+
+        searcher = self.searcher
+        search_alg_state = _find_newest_ckpt(dirpath,
+                                             self.CKPT_FILE_TMPL.format("*"))
+        if not search_alg_state:
+            raise RuntimeError(
+                "Unable to find checkpoint in {}.".format(dirpath))
+        while hasattr(searcher, "searcher"):
+            searcher_name = "name:" + type(searcher).__name__
+            if searcher_name not in search_alg_state:
+                names = [
+                    key.split("name:")[1] for key in search_alg_state
+                    if key.startswith("name:")
+                ]
+                logger.warning("{} was not found in the experiment checkpoint "
+                               "state when restoring. Found {}.".format(
+                                   searcher_name, names))
+            else:
+                searcher.set_state(search_alg_state.pop(searcher_name))
+            searcher = searcher.searcher
+        base_searcher = searcher
+
+        logger.debug(f"searching base {base_searcher}")
+        base_searcher.restore_from_dir(dirpath)
+        self.set_state(search_alg_state)

python/ray/tune/suggest/suggestion.py

@@ -1,23 +1,13 @@
 import copy
+import glob
 import logging
 import os
 
-from ray.tune.error import TuneError
-from ray.tune.experiment import convert_to_experiment_list
-from ray.tune.config_parser import make_parser, create_trial_from_spec
-from ray.tune.suggest.search import SearchAlgorithm
-from ray.tune.suggest.variant_generator import format_vars, resolve_nested_dict
-from ray.tune.trial import Trial
-from ray.tune.utils import merge_dicts, flatten_dict
+from ray.util.debug import log_once
 
 logger = logging.getLogger(__name__)
 
 
-def _warn_on_repeater(searcher, total_samples):
-    from ray.tune.suggest.repeater import _warn_num_samples
-    _warn_num_samples(searcher, total_samples)
-
-
 class Searcher:
     """Abstract class for wrapping suggesting algorithms.
 
@@ -59,7 +49,7 @@ class Searcher:
 
     """
     FINISHED = "FINISHED"
-    CKPT_FILE = "searcher-state.pkl"
+    CKPT_FILE_TMPL = "searcher-state-{}.pkl"
 
     def __init__(self,
                  metric="episode_reward_mean",
@@ -186,23 +176,38 @@ class Searcher:
         """
         raise NotImplementedError
 
-    def save_to_dir(self, checkpoint_dir):
+    def get_state(self):
+        raise NotImplementedError
+
+    def set_state(self, state):
+        raise NotImplementedError
+
+    def save_to_dir(self, checkpoint_dir, session_str="default"):
         """Automatically saves the given searcher to the checkpoint_dir.
 
         This is automatically used by tune.run during a Tune job.
+
+        Args:
+            checkpoint_dir (str): Filepath to experiment dir.
+            session_str (str): Unique identifier of the current run
+                session.
         """
         tmp_search_ckpt_path = os.path.join(checkpoint_dir,
                                             ".tmp_searcher_ckpt")
         success = True
         try:
             self.save(tmp_search_ckpt_path)
-        except NotImplementedError as e:
-            logger.warning(e)
+        except NotImplementedError:
+            with log_once("suggest:save_to_dir"):
+                logger.warning(
+                    "save not implemented for Searcher. Skipping save.")
             success = False
 
         if success and os.path.exists(tmp_search_ckpt_path):
-            os.rename(tmp_search_ckpt_path,
-                      os.path.join(checkpoint_dir, Searcher.CKPT_FILE))
+            os.rename(
+                tmp_search_ckpt_path,
+                os.path.join(checkpoint_dir,
+                             self.CKPT_FILE_TMPL.format(session_str)))
 
     def restore_from_dir(self, checkpoint_dir):
         """Restores the state of a searcher from a given checkpoint_dir.
@@ -225,14 +230,14 @@ class Searcher:
                 os.path.join("~/my_results", self.experiment_name)
         """
 
-        checkpoint_path = os.path.join(checkpoint_dir, Searcher.CKPT_FILE)
-        if os.path.exists(checkpoint_path):
-            self.restore(checkpoint_path)
-        else:
-            raise FileNotFoundError(
-                "{filename} not found in {directory}. Unable to restore "
-                "searcher state from directory.".format(
-                    filename=Searcher.CKPT_FILE, directory=checkpoint_dir))
+        pattern = self.CKPT_FILE_TMPL.format("*")
+        full_paths = glob.glob(os.path.join(checkpoint_dir, pattern))
+        if not full_paths:
+            raise RuntimeError(
+                "Searcher unable to find checkpoint in {}".format(
+                    checkpoint_dir))  # TODO
+        most_recent_checkpoint = max(full_paths)
+        self.restore(most_recent_checkpoint)
 
     @property
     def metric(self):
@@ -271,7 +276,13 @@ class ConcurrencyLimiter(Searcher):
             metric=self.searcher.metric, mode=self.searcher.mode)
 
     def suggest(self, trial_id):
+        assert trial_id not in self.live_trials, (
+            f"Trial ID {trial_id} must be unique: already found in set.")
         if len(self.live_trials) >= self.max_concurrent:
+            logger.debug(
+                f"Not providing a suggestion for {trial_id} due to "
+                "concurrency limit: %s/%s.", len(self.live_trials),
+                self.max_concurrent)
             return
         suggestion = self.searcher.suggest(trial_id)
         if suggestion not in (None, Searcher.FINISHED):
@@ -286,159 +297,10 @@ class ConcurrencyLimiter(Searcher):
                 trial_id, result=result, error=error)
             self.live_trials.remove(trial_id)
 
-    def save(self, checkpoint_dir):
-        self.searcher.save(checkpoint_dir)
+    def get_state(self):
+        state = self.__dict__.copy()
+        del state["searcher"]
+        return copy.deepcopy(state)
 
-    def restore(self, checkpoint_dir):
-        self.searcher.restore(checkpoint_dir)
-
-
-class SearchGenerator(SearchAlgorithm):
-    """Generates trials to be passed to the TrialRunner.
-
-    Uses the provided ``searcher`` object to generate trials. This class
-    transparently handles repeating trials with score aggregation
-    without embedding logic into the Searcher.
-
-    Args:
-        searcher: Search object that subclasses the Searcher base class. This
-            is then used for generating new hyperparameter samples.
-    """
-
-    def __init__(self, searcher):
-        assert issubclass(
-            type(searcher),
-            Searcher), ("Searcher should be subclassing Searcher.")
-        self.searcher = searcher
-        self._parser = make_parser()
-        self._experiment = None
-        self._counter = 0  # Keeps track of number of trials created.
-        self._total_samples = 0  # int: total samples to evaluate.
-        self._finished = False
-
-    def add_configurations(self, experiments):
-        """Registers experiment specifications.
-
-        Arguments:
-            experiments (Experiment | list | dict): Experiments to run.
-        """
-        logger.debug("added configurations")
-        experiment_list = convert_to_experiment_list(experiments)
-        assert len(experiment_list) == 1, (
-            "SearchAlgorithms can only support 1 experiment at a time.")
-        self._experiment = experiment_list[0]
-        experiment_spec = self._experiment.spec
-        self._total_samples = experiment_spec.get("num_samples", 1)
-
-        _warn_on_repeater(self.searcher, self._total_samples)
-
-        if "run" not in experiment_spec:
-            raise TuneError("Must specify `run` in {}".format(experiment_spec))
-
-    def next_trials(self):
-        """Provides a batch of Trial objects to be queued into the TrialRunner.
-
-        Returns:
-            List[Trial]: A list of trials for the Runner to consume.
-        """
-        trials = []
-        while not self.is_finished():
-            trial = self.create_trial_if_possible(self._experiment.spec,
-                                                  self._experiment.name)
-            if trial is None:
-                break
-            trials.append(trial)
-        return trials
-
-    def create_trial_if_possible(self, experiment_spec, output_path):
-        logger.debug("creating trial")
-        trial_id = Trial.generate_id()
-        suggested_config = self.searcher.suggest(trial_id)
-        if suggested_config == Searcher.FINISHED:
-            self._finished = True
-            logger.debug("Searcher has finished.")
-            return
-
-        if suggested_config is None:
-            return
-        spec = copy.deepcopy(experiment_spec)
-        spec["config"] = merge_dicts(spec["config"],
-                                     copy.deepcopy(suggested_config))
-
-        # Create a new trial_id if duplicate trial is created
-        flattened_config = resolve_nested_dict(spec["config"])
-        self._counter += 1
-        tag = "{0}_{1}".format(
-            str(self._counter), format_vars(flattened_config))
-        trial = create_trial_from_spec(
-            spec,
-            output_path,
-            self._parser,
-            evaluated_params=flatten_dict(suggested_config),
-            experiment_tag=tag,
-            trial_id=trial_id)
-        return trial
-
-    def on_trial_result(self, trial_id, result):
-        """Notifies the underlying searcher."""
-        self.searcher.on_trial_result(trial_id, result)
-
-    def on_trial_complete(self, trial_id, result=None, error=False):
-        self.searcher.on_trial_complete(
-            trial_id=trial_id, result=result, error=error)
-
-    def is_finished(self):
-        return self._counter >= self._total_samples or self._finished
-
-    def save(self, checkpoint_path):
-        self.searcher.save(checkpoint_path)
-
-    def restore(self, checkpoint_path):
-        self.searcher.restore(checkpoint_path)
-
-
-class _MockSearcher(Searcher):
-    def __init__(self, **kwargs):
-        self.live_trials = {}
-        self.counter = {"result": 0, "complete": 0}
-        self.final_results = []
-        self.stall = False
-        self.results = []
-        super(_MockSearcher, self).__init__(**kwargs)
-
-    def suggest(self, trial_id):
-        if not self.stall:
-            self.live_trials[trial_id] = 1
-            return {"test_variable": 2}
-        return None
-
-    def on_trial_result(self, trial_id, result):
-        self.counter["result"] += 1
-        self.results += [result]
-
-    def on_trial_complete(self, trial_id, result=None, error=False):
-        self.counter["complete"] += 1
-        if result:
-            self._process_result(result)
-        if trial_id in self.live_trials:
-            del self.live_trials[trial_id]
-
-    def _process_result(self, result):
-        self.final_results += [result]
-
-
-class _MockSuggestionAlgorithm(SearchGenerator):
-    def __init__(self, max_concurrent=None, **kwargs):
-        self.searcher = _MockSearcher(**kwargs)
-        if max_concurrent:
-            self.searcher = ConcurrencyLimiter(
-                self.searcher, max_concurrent=max_concurrent)
-        super(_MockSuggestionAlgorithm, self).__init__(self.searcher)
-
-    @property
-    def live_trials(self):
-        return self.searcher.live_trials
-
-    @property
-    def results(self):
-        return self.searcher.results
+    def set_state(self, state):
+        self.__dict__.update(state)