[tune] refactor tune search space (#10444)

* Added basic functionality and tests * Feature parity with old tune search space config * Convert Optuna search spaces * Introduced quantized values * Updated Optuna resolving * Added HyperOpt search space conversion * Convert search spaces to AxSearch * Convert search spaces to BayesOpt * Added basic functionality and tests * Feature parity with old tune search space config * Convert Optuna search spaces * Introduced quantized values * Updated Optuna resolving * Added HyperOpt search space conversion * Convert search spaces to AxSearch * Convert search spaces to BayesOpt * Re-factored samplers into domain classes * Re-added base classes * Re-factored into list comprehensions * Added `from_config` classmethod for config conversion * Applied suggestions from code review * Removed truncated normal distribution * Set search properties in tune.run * Added test for tune.run search properties * Move sampler initializers to base classes * Add tune API sampling test, fixed includes, fixed resampling bug * Add to API docs * Fix docs * Update metric and mode only when set. Set default metric and mode to experiment analysis object. * Fix experiment analysis tests * Raise error when delimiter is used in the config keys * Added randint/qrandint to API docs, added additional check in tune.run * Fix tests * Fix linting error * Applied suggestions from code review. Re-aded tune.function for the time being * Fix sampling tests * Fix experiment analysis tests * Fix tests and linting error * Removed unnecessary default_config attribute from OptunaSearch * Revert to set AxSearch default metric * fix-min-max * fix * nits * Added function check, enhanced loguniform error message * fix-print * fix * fix * Raise if unresolved values are in config and search space is already set Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
2026-07-18 12:40:56 +08:00 · 2020-09-03 09:06:13 -07:00
parent 715ee8dfc9
commit 06af62ba91
31 changed files with 1548 additions and 211 deletions
@@ -1,3 +1,12 @@
+from typing import Dict
+
+from ax.service.ax_client import AxClient
+from ray.tune.sample import Categorical, Float, Integer, LogUniform, \
+    Quantized, Uniform
+from ray.tune.suggest.variant_generator import parse_spec_vars
+from ray.tune.utils import flatten_dict
+from ray.tune.utils.util import unflatten_dict
+
 try:
    import ax
 except ImportError:
@@ -24,7 +33,7 @@ class AxSearch(Searcher):
        $ pip install ax-platform sqlalchemy

    Parameters:
-        parameters (list[dict]): Parameters in the experiment search space.
+        space (list[dict]): Parameters in the experiment search space.
            Required elements in the dictionaries are: "name" (name of
            this parameter, string), "type" (type of the parameter: "range",
            "fixed", or "choice", string), "bounds" for range parameters
@@ -41,8 +50,11 @@ class AxSearch(Searcher):
            "x3 >= x4" or "x3 + x4 >= 2".
        outcome_constraints (list[str]): Outcome constraints of form
            "metric_name >= bound", like "m1 <= 3."
-        max_concurrent (int): Deprecated.
+        ax_client (AxClient): Optional AxClient instance. If this is set, do
+            not pass any values to these parameters: `space`, `objective_name`,
+            `parameter_constraints`, `outcome_constraints`.
        use_early_stopped_trials: Deprecated.
+        max_concurrent (int): Deprecated.

    .. code-block:: python

@@ -60,41 +72,112 @@ class AxSearch(Searcher):
                intermediate_result = config["x1"] + config["x2"] * i
                tune.report(score=intermediate_result)

-        client = AxClient(enforce_sequential_optimization=False)
-        client.create_experiment(parameters=parameters, objective_name="score")
-        algo = AxSearch(client)
+        client = AxClient()
+        algo = AxSearch(space=parameters, objective_name="score")
        tune.run(easy_objective, search_alg=algo)

    """

    def __init__(self,
-                 ax_client,
+                 space=None,
+                 metric="episode_reward_mean",
                 mode="max",
+                 parameter_constraints=None,
+                 outcome_constraints=None,
+                 ax_client=None,
                 use_early_stopped_trials=None,
                 max_concurrent=None):
        assert ax is not None, "Ax must be installed!"
-        self._ax = ax_client
-        exp = self._ax.experiment
-        self._objective_name = exp.optimization_config.objective.metric.name
-        self.max_concurrent = max_concurrent
-        self._parameters = list(exp.parameters)
-        self._live_trial_mapping = {}
+        assert mode in ["min", "max"], "`mode` must be one of ['min', 'max']"
+
        super(AxSearch, self).__init__(
-            metric=self._objective_name,
+            metric=metric,
            mode=mode,
            max_concurrent=max_concurrent,
            use_early_stopped_trials=use_early_stopped_trials)
+
+        self._ax = ax_client
+        self._space = space
+        self._parameter_constraints = parameter_constraints
+        self._outcome_constraints = outcome_constraints
+
+        self.max_concurrent = max_concurrent
+
+        self._objective_name = metric
+        self._parameters = []
+        self._live_trial_mapping = {}
+
+        if self._ax or self._space:
+            self.setup_experiment()
+
+    def setup_experiment(self):
+        if not self._ax:
+            self._ax = AxClient()
+
+        try:
+            exp = self._ax.experiment
+            has_experiment = True
+        except ValueError:
+            has_experiment = False
+
+        if not has_experiment:
+            if not self._space:
+                raise ValueError(
+                    "You have to create an Ax experiment by calling "
+                    "`AxClient.create_experiment()`, or you should pass an "
+                    "Ax search space as the `space` parameter to `AxSearch`, "
+                    "or pass a `config` dict to `tune.run()`.")
+            self._ax.create_experiment(
+                parameters=self._space,
+                objective_name=self._metric,
+                parameter_constraints=self._parameter_constraints,
+                outcome_constraints=self._outcome_constraints,
+                minimize=self._mode != "max")
+        else:
+            if any([
+                    self._space, self._parameter_constraints,
+                    self._outcome_constraints
+            ]):
+                raise ValueError(
+                    "If you create the Ax experiment yourself, do not pass "
+                    "values for these parameters to `AxSearch`: {}.".format([
+                        "space", "parameter_constraints", "outcome_constraints"
+                    ]))
+
+        exp = self._ax.experiment
+        self._objective_name = exp.optimization_config.objective.metric.name
+        self._parameters = list(exp.parameters)
+
        if self._ax._enforce_sequential_optimization:
            logger.warning("Detected sequential enforcement. Be sure to use "
                           "a ConcurrencyLimiter.")

+    def set_search_properties(self, metric, mode, config):
+        if self._ax:
+            return False
+        space = self.convert_search_space(config)
+        self._space = space
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+        self.setup_experiment()
+        return True
+
    def suggest(self, trial_id):
+        if not self._ax:
+            raise RuntimeError(
+                "Trying to sample a configuration from {}, but no search "
+                "space has been defined. Either pass the `{}` argument when "
+                "instantiating the search algorithm, or pass a `config` to "
+                "`tune.run()`.".format(self.__class__.__name__, "space"))
+
        if self.max_concurrent:
            if len(self._live_trial_mapping) >= self.max_concurrent:
                return None
        parameters, trial_index = self._ax.get_next_trial()
        self._live_trial_mapping[trial_id] = trial_index
-        return parameters
+        return unflatten_dict(parameters)

    def on_trial_complete(self, trial_id, result=None, error=False):
        """Notification for the completion of trial.
@@ -117,3 +200,82 @@ class AxSearch(Searcher):
        metric_dict.update({on: (result[on], 0.0) for on in outcome_names})
        self._ax.complete_trial(
            trial_index=ax_trial_index, raw_data=metric_dict)
+
+    @staticmethod
+    def convert_search_space(spec: Dict):
+        spec = flatten_dict(spec, prevent_delimiter=True)
+        resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
+
+        if grid_vars:
+            raise ValueError(
+                "Grid search parameters cannot be automatically converted "
+                "to an Ax search space.")
+
+        def resolve_value(par, domain):
+            sampler = domain.get_sampler()
+            if isinstance(sampler, Quantized):
+                logger.warning("AxSearch does not support quantization. "
+                               "Dropped quantization.")
+                sampler = sampler.sampler
+
+            if isinstance(domain, Float):
+                if isinstance(sampler, LogUniform):
+                    return {
+                        "name": par,
+                        "type": "range",
+                        "bounds": [domain.lower, domain.upper],
+                        "value_type": "float",
+                        "log_scale": True
+                    }
+                elif isinstance(sampler, Uniform):
+                    return {
+                        "name": par,
+                        "type": "range",
+                        "bounds": [domain.lower, domain.upper],
+                        "value_type": "float",
+                        "log_scale": False
+                    }
+            elif isinstance(domain, Integer):
+                if isinstance(sampler, LogUniform):
+                    return {
+                        "name": par,
+                        "type": "range",
+                        "bounds": [domain.lower, domain.upper],
+                        "value_type": "int",
+                        "log_scale": True
+                    }
+                elif isinstance(sampler, Uniform):
+                    return {
+                        "name": par,
+                        "type": "range",
+                        "bounds": [domain.lower, domain.upper],
+                        "value_type": "int",
+                        "log_scale": False
+                    }
+            elif isinstance(domain, Categorical):
+                if isinstance(sampler, Uniform):
+                    return {
+                        "name": par,
+                        "type": "choice",
+                        "values": domain.categories
+                    }
+
+            raise ValueError("AxSearch does not support parameters of type "
+                             "`{}` with samplers of type `{}`".format(
+                                 type(domain).__name__,
+                                 type(domain.sampler).__name__))
+
+        # Fixed vars
+        fixed_values = [{
+            "name": "/".join(path),
+            "type": "fixed",
+            "value": val
+        } for path, val in resolved_vars]
+
+        # Parameter name is e.g. "a/b/c" for nested dicts
+        resolved_values = [
+            resolve_value("/".join(path), domain)
+            for path, domain in domain_vars
+        ]
+
+        return fixed_values + resolved_values
@@ -1,8 +1,13 @@
-import copy
 from collections import defaultdict
 import logging
 import pickle
 import json
+from typing import Dict
+
+from ray.tune.sample import Float, Quantized
+from ray.tune.suggest.variant_generator import parse_spec_vars
+from ray.tune.utils.util import unflatten_dict
+
 try:  # Python 3 only -- needed for lint test.
    import bayes_opt as byo
 except ImportError:
@@ -76,8 +81,8 @@ class BayesOptSearch(Searcher):
    optimizer = None

    def __init__(self,
-                 space,
-                 metric,
+                 space=None,
+                 metric="episode_reward_mean",
                 mode="max",
                 utility_kwargs=None,
                 random_state=42,
@@ -154,15 +159,45 @@ class BayesOptSearch(Searcher):
        self.random_search_trials = random_search_steps
        self._total_random_search_trials = 0

-        self.optimizer = byo.BayesianOptimization(
-            f=None, pbounds=space, verbose=verbose, random_state=random_state)
-
        self.utility = byo.UtilityFunction(**utility_kwargs)

        # Registering the provided analysis, if given
        if analysis is not None:
            self.register_analysis(analysis)

+        self._space = space
+        self._verbose = verbose
+        self._random_state = random_state
+
+        self.optimizer = None
+        if space:
+            self.setup_optimizer()
+
+    def setup_optimizer(self):
+        self.optimizer = byo.BayesianOptimization(
+            f=None,
+            pbounds=self._space,
+            verbose=self._verbose,
+            random_state=self._random_state)
+
+    def set_search_properties(self, metric, mode, config):
+        if self.optimizer:
+            return False
+        space = self.convert_search_space(config)
+        self._space = space
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+
+        if self._mode == "max":
+            self._metric_op = 1.
+        elif self._mode == "min":
+            self._metric_op = -1.
+
+        self.setup_optimizer()
+        return True
+
    def suggest(self, trial_id):
        """Return new point to be explored by black box function.

@@ -174,6 +209,13 @@ class BayesOptSearch(Searcher):
            Either a dictionary describing the new point to explore or
            None, when no new point is to be explored for the time being.
        """
+        if not self.optimizer:
+            raise RuntimeError(
+                "Trying to sample a configuration from {}, but no search "
+                "space has been defined. Either pass the `{}` argument when "
+                "instantiating the search algorithm, or pass a `config` to "
+                "`tune.run()`.".format(self.__class__.__name__, "space"))
+
        # If we have more active trials than the allowed maximum
        total_live_trials = len(self._live_trial_mapping)
        if self.max_concurrent and self.max_concurrent <= total_live_trials:
@@ -214,7 +256,7 @@ class BayesOptSearch(Searcher):
        self._live_trial_mapping[trial_id] = config

        # Return a deep copy of the mapping
-        return copy.deepcopy(config)
+        return unflatten_dict(config)

    def register_analysis(self, analysis):
        """Integrate the given analysis into the gaussian process.
@@ -283,3 +325,44 @@ class BayesOptSearch(Searcher):
            (self.optimizer, self._buffered_trial_results,
             self._total_random_search_trials,
             self._config_counter) = pickle.load(f)
+
+    @staticmethod
+    def convert_search_space(spec: Dict):
+        spec = flatten_dict(spec, prevent_delimiter=True)
+        resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
+
+        if grid_vars:
+            raise ValueError(
+                "Grid search parameters cannot be automatically converted "
+                "to a BayesOpt search space.")
+
+        if resolved_vars:
+            raise ValueError(
+                "BayesOpt does not support fixed parameters. Please find a "
+                "different way to pass constants to your training function.")
+
+        def resolve_value(domain):
+            sampler = domain.get_sampler()
+            if isinstance(sampler, Quantized):
+                logger.warning(
+                    "BayesOpt search does not support quantization. "
+                    "Dropped quantization.")
+                sampler = sampler.get_sampler()
+
+            if isinstance(domain, Float):
+                if domain.sampler is not None:
+                    logger.warning(
+                        "BayesOpt does not support specific sampling methods. "
+                        "The {} sampler will be dropped.".format(sampler))
+                    return (domain.lower, domain.upper)
+
+            raise ValueError("BayesOpt does not support parameters of type "
+                             "`{}`".format(type(domain).__name__))
+
+        # Parameter name is e.g. "a/b/c" for nested dicts
+        bounds = {
+            "/".join(path): resolve_value(domain)
+            for path, domain in domain_vars
+        }
+
+        return bounds
@@ -1,8 +1,16 @@
+from typing import Dict
+
 import numpy as np
 import copy
 import logging
 from functools import partial
 import pickle
+
+from ray.tune.sample import Categorical, Float, Integer, LogUniform, Normal, \
+    Quantized, \
+    Uniform
+from ray.tune.suggest.variant_generator import assign_value, parse_spec_vars
+
 try:
    hyperopt_logger = logging.getLogger("hyperopt")
    hyperopt_logger.setLevel(logging.WARNING)
@@ -84,7 +92,7 @@ class HyperOptSearch(Searcher):

    def __init__(
            self,
-            space,
+            space=None,
            metric="episode_reward_mean",
            mode="max",
            points_to_evaluate=None,
@@ -116,7 +124,6 @@ class HyperOptSearch(Searcher):
                hpo.tpe.suggest, n_startup_jobs=n_initial_points)
        if gamma is not None:
            self.algo = partial(self.algo, gamma=gamma)
-        self.domain = hpo.Domain(lambda spc: spc, space)
        if points_to_evaluate is None:
            self._hpopt_trials = hpo.Trials()
            self._points_to_evaluate = 0
@@ -132,7 +139,35 @@ class HyperOptSearch(Searcher):
        else:
            self.rstate = np.random.RandomState(random_state_seed)

+        self.domain = None
+        if space:
+            self.domain = hpo.Domain(lambda spc: spc, space)
+
+    def set_search_properties(self, metric, mode, config):
+        if self.domain:
+            return False
+        space = self.convert_search_space(config)
+        self.domain = hpo.Domain(lambda spc: spc, space)
+
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+
+        if self._mode == "max":
+            self.metric_op = -1.
+        elif self._mode == "min":
+            self.metric_op = 1.
+
+        return True
+
    def suggest(self, trial_id):
+        if not self.domain:
+            raise RuntimeError(
+                "Trying to sample a configuration from {}, but no search "
+                "space has been defined. Either pass the `{}` argument when "
+                "instantiating the search algorithm, or pass a `config` to "
+                "`tune.run()`.".format(self.__class__.__name__, "space"))
        if self.max_concurrent:
            if len(self._live_trial_mapping) >= self.max_concurrent:
                return None
@@ -235,3 +270,75 @@ class HyperOptSearch(Searcher):
            self.rstate.set_state(trials_object[1])
        else:
            self.set_state(trials_object)
+
+    @staticmethod
+    def convert_search_space(spec: Dict):
+        spec = copy.deepcopy(spec)
+        resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
+
+        if not domain_vars and not grid_vars:
+            return []
+
+        if grid_vars:
+            raise ValueError(
+                "Grid search parameters cannot be automatically converted "
+                "to a HyperOpt search space.")
+
+        def resolve_value(par, domain):
+            quantize = None
+
+            sampler = domain.get_sampler()
+            if isinstance(sampler, Quantized):
+                quantize = sampler.q
+                sampler = sampler.sampler
+
+            if isinstance(domain, Float):
+                if isinstance(sampler, LogUniform):
+                    if quantize:
+                        return hpo.hp.qloguniform(par, domain.lower,
+                                                  domain.upper, quantize)
+                    return hpo.hp.loguniform(par, np.log(domain.lower),
+                                             np.log(domain.upper))
+                elif isinstance(sampler, Uniform):
+                    if quantize:
+                        return hpo.hp.quniform(par, domain.lower, domain.upper,
+                                               quantize)
+                    return hpo.hp.uniform(par, domain.lower, domain.upper)
+                elif isinstance(sampler, Normal):
+                    if quantize:
+                        return hpo.hp.qnormal(par, sampler.mean, sampler.sd,
+                                              quantize)
+                    return hpo.hp.normal(par, sampler.mean, sampler.sd)
+
+            elif isinstance(domain, Integer):
+                if isinstance(sampler, Uniform):
+                    if quantize:
+                        logger.warning(
+                            "HyperOpt does not support quantization for "
+                            "integer values. Reverting back to 'randint'.")
+                    if domain.lower != 0:
+                        raise ValueError(
+                            "HyperOpt only allows integer sampling with "
+                            f"lower bound 0. Got: {domain.lower}.")
+                    if domain.upper < 1:
+                        raise ValueError(
+                            "HyperOpt does not support integer sampling "
+                            "of values lower than 0. Set your maximum range "
+                            "to something above 0 (currently {})".format(
+                                domain.upper))
+                    return hpo.hp.randint(par, domain.upper)
+            elif isinstance(domain, Categorical):
+                if isinstance(sampler, Uniform):
+                    return hpo.hp.choice(par, domain.categories)
+
+            raise ValueError("HyperOpt does not support parameters of type "
+                             "`{}` with samplers of type `{}`".format(
+                                 type(domain).__name__,
+                                 type(domain.sampler).__name__))
+
+        for path, domain in domain_vars:
+            par = "/".join(path)
+            value = resolve_value(par, domain)
+            assign_value(spec, path, value)
+
+        return spec
@@ -115,7 +115,6 @@ class NevergradSearch(Searcher):
        # in v0.2.0+, output of ask() is a Candidate,
        # with fields args and kwargs
        if not suggested_config.kwargs:
-            print(suggested_config.args, suggested_config.kwargs)
            return dict(zip(self._parameters, suggested_config.args[0]))
        else:
            return suggested_config.kwargs
@@ -1,7 +1,13 @@
 import logging
 import pickle
+from typing import Dict

 from ray.tune.result import TRAINING_ITERATION
+from ray.tune.sample import Categorical, Float, Integer, LogUniform, \
+    Quantized, Uniform
+from ray.tune.suggest.variant_generator import parse_spec_vars
+from ray.tune.utils import flatten_dict
+from ray.tune.utils.util import unflatten_dict

 try:
    import optuna as ot
@@ -74,13 +80,11 @@ class OptunaSearch(Searcher):

    """

-    def __init__(
-            self,
-            space,
-            metric="episode_reward_mean",
-            mode="max",
-            sampler=None,
-    ):
+    def __init__(self,
+                 space=None,
+                 metric="episode_reward_mean",
+                 mode="max",
+                 sampler=None):
        assert ot is not None, (
            "Optuna must be installed! Run `pip install optuna`.")
        super(OptunaSearch, self).__init__(
@@ -101,6 +105,11 @@ class OptunaSearch(Searcher):
        self._storage = ot.storages.InMemoryStorage()

        self._ot_trials = {}
+        self._ot_study = None
+        if self._space:
+            self.setup_study(mode)
+
+    def setup_study(self, mode):
        self._ot_study = ot.study.create_study(
            storage=self._storage,
            sampler=self._sampler,
@@ -109,18 +118,40 @@ class OptunaSearch(Searcher):
            direction="minimize" if mode == "min" else "maximize",
            load_if_exists=True)

+    def set_search_properties(self, metric, mode, config):
+        if self._space:
+            return False
+        space = self.convert_search_space(config)
+        self._space = space
+        if metric:
+            self._metric = metric
+        if mode:
+            self._mode = mode
+        self.setup_study(mode)
+        return True
+
    def suggest(self, trial_id):
+        if not self._space:
+            raise RuntimeError(
+                "Trying to sample a configuration from {}, but no search "
+                "space has been defined. Either pass the `{}` argument when "
+                "instantiating the search algorithm, or pass a `config` to "
+                "`tune.run()`.".format(self.__class__.__name__, "space"))
+
        if trial_id not in self._ot_trials:
            ot_trial_id = self._storage.create_new_trial(
                self._ot_study._study_id)
            self._ot_trials[trial_id] = ot.trial.Trial(self._ot_study,
                                                       ot_trial_id)
        ot_trial = self._ot_trials[trial_id]
-        params = {}
-        for (fn, args, kwargs) in self._space:
-            param_name = args[0] if len(args) > 0 else kwargs["name"]
-            params[param_name] = getattr(ot_trial, fn)(*args, **kwargs)
-        return params
+
+        # getattr will fetch the trial.suggest_ function on Optuna trials
+        params = {
+            args[0] if len(args) > 0 else kwargs["name"]: getattr(
+                ot_trial, fn)(*args, **kwargs)
+            for (fn, args, kwargs) in self._space
+        }
+        return unflatten_dict(params)

    def on_trial_result(self, trial_id, result):
        metric = result[self.metric]
@@ -147,3 +178,67 @@ class OptunaSearch(Searcher):
            save_object = pickle.load(inputFile)
        self._storage, self._pruner, self._sampler, \
            self._ot_trials, self._ot_study = save_object
+
+    @staticmethod
+    def convert_search_space(spec: Dict):
+        spec = flatten_dict(spec, prevent_delimiter=True)
+        resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
+
+        if not domain_vars and not grid_vars:
+            return []
+
+        if grid_vars:
+            raise ValueError(
+                "Grid search parameters cannot be automatically converted "
+                "to an Optuna search space.")
+
+        def resolve_value(par, domain):
+            quantize = None
+
+            sampler = domain.get_sampler()
+            if isinstance(sampler, Quantized):
+                quantize = sampler.q
+                sampler = sampler.sampler
+
+            if isinstance(domain, Float):
+                if isinstance(sampler, LogUniform):
+                    if quantize:
+                        logger.warning(
+                            "Optuna does not support both quantization and "
+                            "sampling from LogUniform. Dropped quantization.")
+                    return param.suggest_loguniform(par, domain.lower,
+                                                    domain.upper)
+                elif isinstance(sampler, Uniform):
+                    if quantize:
+                        return param.suggest_discrete_uniform(
+                            par, domain.lower, domain.upper, quantize)
+                    return param.suggest_uniform(par, domain.lower,
+                                                 domain.upper)
+            elif isinstance(domain, Integer):
+                if isinstance(sampler, LogUniform):
+                    if quantize:
+                        logger.warning(
+                            "Optuna does not support both quantization and "
+                            "sampling from LogUniform. Dropped quantization.")
+                    return param.suggest_int(
+                        par, domain.lower, domain.upper, log=True)
+                elif isinstance(sampler, Uniform):
+                    return param.suggest_int(
+                        par, domain.lower, domain.upper, step=quantize or 1)
+            elif isinstance(domain, Categorical):
+                if isinstance(sampler, Uniform):
+                    return param.suggest_categorical(par, domain.categories)
+
+            raise ValueError(
+                "Optuna search does not support parameters of type "
+                "`{}` with samplers of type `{}`".format(
+                    type(domain).__name__,
+                    type(domain.sampler).__name__))
+
+        # Parameter name is e.g. "a/b/c" for nested dicts
+        values = [
+            resolve_value("/".join(path), domain)
+            for path, domain in domain_vars
+        ]
+
+        return values
@@ -12,6 +12,23 @@ class SearchAlgorithm:
    """
    _finished = False

+    def set_search_properties(self, metric, mode, config):
+        """Pass search properties to search algorithm.
+
+        This method acts as an alternative to instantiating search algorithms
+        with their own specific search spaces. Instead they can accept a
+        Tune config through this method.
+
+        The search algorithm will usually pass this method to their
+        ``Searcher`` instance.
+
+        Args:
+            metric (str): Metric to optimize
+            mode (str): One of ["min", "max"]. Direction to optimize.
+            config (dict): Tune config dict.
+        """
+        return True
+
    def add_configurations(self, experiments):
        """Tracks given experiment specifications.

@@ -69,6 +69,9 @@ class SearchGenerator(SearchAlgorithm):
        self._total_samples = None  # int: total samples to evaluate.
        self._finished = False

+    def set_search_properties(self, metric, mode, config):
+        return self.searcher.set_search_properties(metric, mode, config)
+
    def add_configurations(self, experiments):
        """Registers experiment specifications.

@@ -86,6 +86,22 @@ class Searcher:
        self._metric = metric
        self._mode = mode

+    def set_search_properties(self, metric, mode, config):
+        """Pass search properties to searcher.
+
+        This method acts as an alternative to instantiating search algorithms
+        with their own specific search spaces. Instead they can accept a
+        Tune config through this method. A searcher should return ``True``
+        if setting the config was successful, or ``False`` if it was
+        unsuccessful, e.g. when the search space has already been set.
+
+        Args:
+            metric (str): Metric to optimize
+            mode (str): One of ["min", "max"]. Direction to optimize.
+            config (dict): Tune config dict.
+        """
+        return False
+
    def on_trial_result(self, trial_id, result):
        """Optional notification for result during training.

@@ -4,7 +4,7 @@ import numpy
 import random

 from ray.tune import TuneError
-from ray.tune.sample import sample_from
+from ray.tune.sample import Categorical, Domain, Function

 logger = logging.getLogger(__name__)

@@ -115,25 +115,36 @@ def _clean_value(value):
        return str(value).replace("/", "_")


+def parse_spec_vars(spec):
+    resolved, unresolved = _split_resolved_unresolved_values(spec)
+    resolved_vars = list(resolved.items())
+
+    if not unresolved:
+        return resolved_vars, [], []
+
+    grid_vars = []
+    domain_vars = []
+    for path, value in unresolved.items():
+        if value.is_grid():
+            grid_vars.append((path, value))
+        else:
+            domain_vars.append((path, value))
+    grid_vars.sort()
+
+    return resolved_vars, domain_vars, grid_vars
+
+
 def _generate_variants(spec):
    spec = copy.deepcopy(spec)
-    unresolved = _unresolved_values(spec)
-    if not unresolved:
+    _, domain_vars, grid_vars = parse_spec_vars(spec)
+
+    if not domain_vars and not grid_vars:
        yield {}, spec
        return

-    grid_vars = []
-    lambda_vars = []
-    for path, value in unresolved.items():
-        if callable(value):
-            lambda_vars.append((path, value))
-        else:
-            grid_vars.append((path, value))
-    grid_vars.sort()
-
    grid_search = _grid_search_generator(spec, grid_vars)
    for resolved_spec in grid_search:
-        resolved_vars = _resolve_lambda_vars(resolved_spec, lambda_vars)
+        resolved_vars = _resolve_domain_vars(resolved_spec, domain_vars)
        for resolved, spec in _generate_variants(resolved_spec):
            for path, value in grid_vars:
                resolved_vars[path] = _get_value(spec, path)
@@ -148,7 +159,7 @@ def _generate_variants(spec):
            yield resolved_vars, spec


-def _assign_value(spec, path, value):
+def assign_value(spec, path, value):
    for k in path[:-1]:
        spec = spec[k]
    spec[path[-1]] = value
@@ -160,23 +171,26 @@ def _get_value(spec, path):
    return spec


-def _resolve_lambda_vars(spec, lambda_vars):
+def _resolve_domain_vars(spec, domain_vars):
    resolved = {}
    error = True
    num_passes = 0
    while error and num_passes < _MAX_RESOLUTION_PASSES:
        num_passes += 1
        error = False
-        for path, fn in lambda_vars:
+        for path, domain in domain_vars:
+            if path in resolved:
+                continue
            try:
-                value = fn(_UnresolvedAccessGuard(spec))
+                value = domain.sample(_UnresolvedAccessGuard(spec))
            except RecursiveDependencyError as e:
                error = e
            except Exception:
                raise ValueError(
-                    "Failed to evaluate expression: {}: {}".format(path, fn))
+                    "Failed to evaluate expression: {}: {}".format(
+                        path, domain))
            else:
-                _assign_value(spec, path, value)
+                assign_value(spec, path, value)
                resolved[path] = value
    if error:
        raise error
@@ -203,7 +217,7 @@ def _grid_search_generator(unresolved_spec, grid_vars):
    while value_indices[-1] < len(grid_vars[-1][1]):
        spec = copy.deepcopy(unresolved_spec)
        for i, (path, values) in enumerate(grid_vars):
-            _assign_value(spec, path, values[value_indices[i]])
+            assign_value(spec, path, values[value_indices[i]])
        yield spec
        if grid_vars:
            done = increment(0)
@@ -217,13 +231,13 @@ def _is_resolved(v):


 def _try_resolve(v):
-    if isinstance(v, sample_from):
-        # Function to sample from
-        return False, v.func
+    if isinstance(v, Domain):
+        # Domain to sample from
+        return False, v
    elif isinstance(v, dict) and len(v) == 1 and "eval" in v:
        # Lambda function in eval syntax
-        return False, lambda spec: eval(
-            v["eval"], _STANDARD_IMPORTS, {"spec": spec})
+        return False, Function(
+            lambda spec: eval(v["eval"], _STANDARD_IMPORTS, {"spec": spec}))
    elif isinstance(v, dict) and len(v) == 1 and "grid_search" in v:
        # Grid search values
        grid_values = v["grid_search"]
@@ -231,26 +245,45 @@ def _try_resolve(v):
            raise TuneError(
                "Grid search expected list of values, got: {}".format(
                    grid_values))
-        return False, grid_values
+        return False, Categorical(grid_values).grid()
    return True, v


-def _unresolved_values(spec):
-    found = {}
+def _split_resolved_unresolved_values(spec):
+    resolved_vars = {}
+    unresolved_vars = {}
    for k, v in spec.items():
        resolved, v = _try_resolve(v)
        if not resolved:
-            found[(k, )] = v
+            unresolved_vars[(k, )] = v
        elif isinstance(v, dict):
            # Recurse into a dict
-            for (path, value) in _unresolved_values(v).items():
-                found[(k, ) + path] = value
+            _resolved_children, _unresolved_children = \
+                _split_resolved_unresolved_values(v)
+            for (path, value) in _resolved_children.items():
+                resolved_vars[(k, ) + path] = value
+            for (path, value) in _unresolved_children.items():
+                unresolved_vars[(k, ) + path] = value
        elif isinstance(v, list):
            # Recurse into a list
            for i, elem in enumerate(v):
-                for (path, value) in _unresolved_values({i: elem}).items():
-                    found[(k, ) + path] = value
-    return found
+                _resolved_children, _unresolved_children = \
+                    _split_resolved_unresolved_values({i: elem})
+                for (path, value) in _resolved_children.items():
+                    resolved_vars[(k, ) + path] = value
+                for (path, value) in _unresolved_children.items():
+                    unresolved_vars[(k, ) + path] = value
+        else:
+            resolved_vars[(k, )] = v
+    return resolved_vars, unresolved_vars
+
+
+def _unresolved_values(spec):
+    return _split_resolved_unresolved_values(spec)[1]
+
+
+def has_unresolved_values(spec):
+    return True if _unresolved_values(spec) else False


 class _UnresolvedAccessGuard(dict):