autogen/flaml/searcher/blendsearch.py

# !
#  * Copyright (c) Microsoft Corporation. All rights reserved.
#  * Licensed under the MIT License. See LICENSE file in the
#  * project root for license information.
from typing import Dict, Optional, List, Tuple, Callable, Union
import numpy as np
import time
import pickle


try:
    from ray import __version__ as ray_version

    assert ray_version >= "1.0.0"
    from ray.tune.suggest import Searcher
    from ray.tune.suggest.optuna import OptunaSearch as GlobalSearch
except (ImportError, AssertionError):
    from .suggestion import Searcher
    from .suggestion import OptunaSearch as GlobalSearch
from ..tune.trial import unflatten_dict, flatten_dict
from .search_thread import SearchThread
from .flow2 import FLOW2
from ..tune.space import add_cost_to_space, indexof, normalize, define_by_run_func
import logging

logger = logging.getLogger(__name__)


class BlendSearch(Searcher):
    """class for BlendSearch algorithm."""

    cost_attr = "time_total_s"  # cost attribute in result
    lagrange = "_lagrange"  # suffix for lagrange-modified metric
    penalty = 1e10  # penalty term for constraints
    LocalSearch = FLOW2

    def __init__(
        self,
        metric: Optional[str] = None,
        mode: Optional[str] = None,
        space: Optional[dict] = None,
        low_cost_partial_config: Optional[dict] = None,
        cat_hp_cost: Optional[dict] = None,
        points_to_evaluate: Optional[List[dict]] = None,
        evaluated_rewards: Optional[List] = None,
        time_budget_s: Union[int, float] = None,
        num_samples: Optional[int] = None,
        prune_attr: Optional[str] = None,
        min_resource: Optional[float] = None,
        max_resource: Optional[float] = None,
        reduction_factor: Optional[float] = None,
        global_search_alg: Optional[Searcher] = None,
        config_constraints: Optional[
            List[Tuple[Callable[[dict], float], str, float]]
        ] = None,
        metric_constraints: Optional[List[Tuple[str, str, float]]] = None,
        seed: Optional[int] = 20,
        experimental: Optional[bool] = False,
    ):
        """Constructor.

        Args:
            metric: A string of the metric name to optimize for.
            mode: A string in ['min', 'max'] to specify the objective as
                minimization or maximization.
            space: A dictionary to specify the search space.
            low_cost_partial_config: A dictionary from a subset of
                controlled dimensions to the initial low-cost values.
                e.g.,

                .. code-block:: python

                    {'n_estimators': 4, 'max_leaves': 4}

            cat_hp_cost: A dictionary from a subset of categorical dimensions
                to the relative cost of each choice.
                e.g.,

                .. code-block:: python

                    {'tree_method': [1, 1, 2]}

                i.e., the relative cost of the
                three choices of 'tree_method' is 1, 1 and 2 respectively.
            points_to_evaluate: Initial parameter suggestions to be run first.
            evaluated_rewards (list): If you have previously evaluated the
                parameters passed in as points_to_evaluate you can avoid
                re-running those trials by passing in the reward attributes
                as a list so the optimiser can be told the results without
                needing to re-compute the trial. Must be the same length as
                points_to_evaluate.
            time_budget_s: int or float | Time budget in seconds.
            num_samples: int | The number of configs to try.
            prune_attr: A string of the attribute used for pruning.
                Not necessarily in space.
                When prune_attr is in space, it is a hyperparameter, e.g.,
                    'n_iters', and the best value is unknown.
                When prune_attr is not in space, it is a resource dimension,
                    e.g., 'sample_size', and the peak performance is assumed
                    to be at the max_resource.
            min_resource: A float of the minimal resource to use for the
                prune_attr; only valid if prune_attr is not in space.
            max_resource: A float of the maximal resource to use for the
                prune_attr; only valid if prune_attr is not in space.
            reduction_factor: A float of the reduction factor used for
                incremental pruning.
            global_search_alg: A Searcher instance as the global search
                instance. If omitted, Optuna is used. The following algos have
                known issues when used as global_search_alg:
                - HyperOptSearch raises exception sometimes
                - TuneBOHB has its own scheduler
            config_constraints: A list of config constraints to be satisfied.
                e.g.,

                .. code-block: python

                    config_constraints = [(mem_size, '<=', 1024**3)]

                mem_size is a function which produces a float number for the bytes
                needed for a config.
                It is used to skip configs which do not fit in memory.
            metric_constraints: A list of metric constraints to be satisfied.
                e.g., `['precision', '>=', 0.9]`
            seed: An integer of the random seed.
            experimental: A bool of whether to use experimental features.
        """
        self._metric, self._mode = metric, mode
        init_config = low_cost_partial_config or {}
        if not init_config:
            logger.info(
                "No low-cost partial config given to the search algorithm. "
                "For cost-frugal search, "
                "consider providing low-cost values for cost-related hps via "
                "'low_cost_partial_config'. More info can be found at "
                "https://github.com/microsoft/FLAML/wiki/About-%60low_cost_partial_config%60"
            )
        if evaluated_rewards and mode:
            self._points_to_evaluate = []
            self._evaluated_rewards = []
            best = max(evaluated_rewards) if mode == "max" else min(evaluated_rewards)
            # only keep the best points as start points
            for i, r in enumerate(evaluated_rewards):
                if r == best:
                    p = points_to_evaluate[i]
                    self._points_to_evaluate.append(p)
                    self._evaluated_rewards.append(r)
        else:
            self._points_to_evaluate = points_to_evaluate or []
            self._evaluated_rewards = evaluated_rewards or []
        self._config_constraints = config_constraints
        self._metric_constraints = metric_constraints
        if self._metric_constraints:
            # metric modified by lagrange
            metric += self.lagrange
        self._cat_hp_cost = cat_hp_cost or {}
        if space:
            add_cost_to_space(space, init_config, self._cat_hp_cost)
        self._ls = self.LocalSearch(
            init_config,
            metric,
            mode,
            space,
            prune_attr,
            min_resource,
            max_resource,
            reduction_factor,
            self.cost_attr,
            seed,
        )
        if global_search_alg is not None:
            self._gs = global_search_alg
        elif getattr(self, "__name__", None) != "CFO":
            if space and self._ls.hierarchical:
                from functools import partial

                gs_space = partial(define_by_run_func, space=space)
                evaluated_rewards = None  # not supproted by define-by-run
            else:
                gs_space = space
            gs_seed = seed - 10 if (seed - 10) >= 0 else seed - 11 + (1 << 32)
            if experimental:
                import optuna as ot

                sampler = ot.samplers.TPESampler(
                    seed=seed, multivariate=True, group=True
                )
            else:
                sampler = None
            try:
                self._gs = GlobalSearch(
                    space=gs_space,
                    metric=metric,
                    mode=mode,
                    seed=gs_seed,
                    sampler=sampler,
                    points_to_evaluate=points_to_evaluate,
                    evaluated_rewards=evaluated_rewards,
                )
            except ValueError:
                self._gs = GlobalSearch(
                    space=gs_space,
                    metric=metric,
                    mode=mode,
                    seed=gs_seed,
                    sampler=sampler,
                )
            self._gs.space = space
        else:
            self._gs = None
        self._experimental = experimental
        if (
            getattr(self, "__name__", None) == "CFO"
            and points_to_evaluate
            and len(self._points_to_evaluate) > 1
        ):
            # use the best config in points_to_evaluate as the start point
            self._candidate_start_points = {}
            self._started_from_low_cost = not low_cost_partial_config
        else:
            self._candidate_start_points = None
        self._time_budget_s, self._num_samples = time_budget_s, num_samples
        if space is not None:
            self._init_search()

    def set_search_properties(
        self,
        metric: Optional[str] = None,
        mode: Optional[str] = None,
        config: Optional[Dict] = None,
        setting: Optional[Dict] = None,
    ) -> bool:
        metric_changed = mode_changed = False
        if metric and self._metric != metric:
            metric_changed = True
            self._metric = metric
            if self._metric_constraints:
                # metric modified by lagrange
                metric += self.lagrange
                # TODO: don't change metric for global search methods that
                # can handle constraints already
        if mode and self._mode != mode:
            mode_changed = True
            self._mode = mode
        if not self._ls.space:
            # the search space can be set only once
            if self._gs is not None:
                # define-by-run is not supported via set_search_properties
                self._gs.set_search_properties(metric, mode, config)
                self._gs.space = config
            if config:
                add_cost_to_space(config, self._ls.init_config, self._cat_hp_cost)
            self._ls.set_search_properties(metric, mode, config)
            self._init_search()
        else:
            if metric_changed or mode_changed:
                # reset search when metric or mode changed
                self._ls.set_search_properties(metric, mode)
                if self._gs is not None:
                    self._gs = GlobalSearch(
                        space=self._gs._space,
                        metric=metric,
                        mode=mode,
                        sampler=self._gs._sampler,
                    )
                    self._gs.space = self._ls.space
                self._init_search()
        if setting:
            # CFO doesn't need these settings
            if "time_budget_s" in setting:
                self._time_budget_s = setting["time_budget_s"]  # budget from now
                now = time.time()
                self._time_used += now - self._start_time
                self._start_time = now
                self._set_deadline()
            if "metric_target" in setting:
                self._metric_target = setting.get("metric_target")
            if "num_samples" in setting:
                self._num_samples = (
                    setting["num_samples"]
                    + len(self._result)
                    + len(self._trial_proposed_by)
                )
        return True

    def _set_deadline(self):
        if self._time_budget_s is not None:
            self._deadline = self._time_budget_s + self._start_time
            SearchThread.set_eps(self._time_budget_s)
        else:
            self._deadline = np.inf

    def _init_search(self):
        """initialize the search"""
        self._start_time = time.time()
        self._time_used = 0
        self._set_deadline()
        self._is_ls_ever_converged = False
        self._subspace = {}  # the subspace for each trial id
        self._metric_target = np.inf * self._ls.metric_op
        self._search_thread_pool = {
            # id: int -> thread: SearchThread
            0: SearchThread(self._ls.mode, self._gs)
        }
        self._thread_count = 1  # total # threads created
        self._init_used = self._ls.init_config is None
        self._trial_proposed_by = {}  # trial_id: str -> thread_id: int
        self._ls_bound_min = normalize(
            self._ls.init_config.copy(),
            self._ls.space,
            self._ls.init_config,
            {},
            recursive=True,
        )
        self._ls_bound_max = normalize(
            self._ls.init_config.copy(),
            self._ls.space,
            self._ls.init_config,
            {},
            recursive=True,
        )
        self._gs_admissible_min = self._ls_bound_min.copy()
        self._gs_admissible_max = self._ls_bound_max.copy()
        self._result = {}  # config_signature: tuple -> result: Dict
        if self._metric_constraints:
            self._metric_constraint_satisfied = False
            self._metric_constraint_penalty = [
                self.penalty for _ in self._metric_constraints
            ]
        else:
            self._metric_constraint_satisfied = True
            self._metric_constraint_penalty = None
        self.best_resource = self._ls.min_resource

    def save(self, checkpoint_path: str):
        """save states to a checkpoint path."""
        self._time_used += time.time() - self._start_time
        self._start_time = time.time()
        save_object = self
        with open(checkpoint_path, "wb") as outputFile:
            pickle.dump(save_object, outputFile)

    def restore(self, checkpoint_path: str):
        """restore states from checkpoint."""
        with open(checkpoint_path, "rb") as inputFile:
            state = pickle.load(inputFile)
        self.__dict__ = state.__dict__
        self._start_time = time.time()
        self._set_deadline()

    @property
    def metric_target(self):
        return self._metric_target

    @property
    def is_ls_ever_converged(self):
        return self._is_ls_ever_converged

    def on_trial_complete(
        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
    ):
        """search thread updater and cleaner."""
        metric_constraint_satisfied = True
        if result and not error and self._metric_constraints:
            # account for metric constraints if any
            objective = result[self._metric]
            for i, constraint in enumerate(self._metric_constraints):
                metric_constraint, sign, threshold = constraint
                value = result.get(metric_constraint)
                if value:
                    # sign is <= or >=
                    sign_op = 1 if sign == "<=" else -1
                    violation = (value - threshold) * sign_op
                    if violation > 0:
                        # add penalty term to the metric
                        objective += (
                            self._metric_constraint_penalty[i]
                            * violation
                            * self._ls.metric_op
                        )
                        metric_constraint_satisfied = False
                        if self._metric_constraint_penalty[i] < self.penalty:
                            self._metric_constraint_penalty[i] += violation
            result[self._metric + self.lagrange] = objective
            if metric_constraint_satisfied and not self._metric_constraint_satisfied:
                # found a feasible point
                self._metric_constraint_penalty = [1 for _ in self._metric_constraints]
            self._metric_constraint_satisfied |= metric_constraint_satisfied
        thread_id = self._trial_proposed_by.get(trial_id)
        if thread_id in self._search_thread_pool:
            self._search_thread_pool[thread_id].on_trial_complete(
                trial_id, result, error
            )
            del self._trial_proposed_by[trial_id]
        if result:
            config = result.get("config", {})
            if not config:
                for key, value in result.items():
                    if key.startswith("config/"):
                        config[key[7:]] = value
            signature = self._ls.config_signature(
                config, self._subspace.get(trial_id, {})
            )
            if error:  # remove from result cache
                del self._result[signature]
            else:  # add to result cache
                self._result[signature] = result
                # update target metric if improved
                objective = result[self._ls.metric]
                if (objective - self._metric_target) * self._ls.metric_op < 0:
                    self._metric_target = objective
                    if self._ls.resource:
                        self._best_resource = config[self._ls.prune_attr]
                if thread_id:
                    if not self._metric_constraint_satisfied:
                        # no point has been found to satisfy metric constraint
                        self._expand_admissible_region(
                            self._ls_bound_min,
                            self._ls_bound_max,
                            self._subspace.get(trial_id, self._ls.space),
                        )
                    if (
                        self._gs is not None
                        and self._experimental
                        and (not self._ls.hierarchical)
                    ):
                        self._gs.add_evaluated_point(flatten_dict(config), objective)
                        # TODO: recover when supported
                        # converted = convert_key(config, self._gs.space)
                        # logger.info(converted)
                        # self._gs.add_evaluated_point(converted, objective)
                elif metric_constraint_satisfied and self._create_condition(result):
                    # thread creator
                    thread_id = self._thread_count
                    self._started_from_given = (
                        self._candidate_start_points
                        and trial_id in self._candidate_start_points
                    )
                    if self._started_from_given:
                        del self._candidate_start_points[trial_id]
                    else:
                        self._started_from_low_cost = True
                    self._create_thread(
                        config, result, self._subspace.get(trial_id, self._ls.space)
                    )
                # reset admissible region to ls bounding box
                self._gs_admissible_min.update(self._ls_bound_min)
                self._gs_admissible_max.update(self._ls_bound_max)
        # cleaner
        if thread_id and thread_id in self._search_thread_pool:
            # local search thread
            self._clean(thread_id)
        if trial_id in self._subspace and not (
            self._candidate_start_points and trial_id in self._candidate_start_points
        ):
            del self._subspace[trial_id]

    def _create_thread(self, config, result, space):
        self._search_thread_pool[self._thread_count] = SearchThread(
            self._ls.mode,
            self._ls.create(
                config,
                result[self._ls.metric],
                cost=result.get(self.cost_attr, 1),
                space=space,
            ),
            self.cost_attr,
        )
        self._thread_count += 1
        self._update_admissible_region(
            unflatten_dict(config),
            self._ls_bound_min,
            self._ls_bound_max,
            space,
            self._ls.space,
        )

    def _update_admissible_region(
        self,
        config,
        admissible_min,
        admissible_max,
        subspace: Dict = {},
        space: Dict = {},
    ):
        # update admissible region
        normalized_config = normalize(config, subspace, config, {})
        for key in admissible_min:
            value = normalized_config[key]
            if isinstance(admissible_max[key], list):
                domain = space[key]
                choice = indexof(domain, value)
                self._update_admissible_region(
                    value,
                    admissible_min[key][choice],
                    admissible_max[key][choice],
                    subspace[key],
                    domain[choice],
                )
                if len(admissible_max[key]) > len(domain.categories):
                    # points + index
                    normal = (choice + 0.5) / len(domain.categories)
                    admissible_max[key][-1] = max(normal, admissible_max[key][-1])
                    admissible_min[key][-1] = min(normal, admissible_min[key][-1])
            elif isinstance(value, dict):
                self._update_admissible_region(
                    value,
                    admissible_min[key],
                    admissible_max[key],
                    subspace[key],
                    space[key],
                )
            else:
                if value > admissible_max[key]:
                    admissible_max[key] = value
                elif value < admissible_min[key]:
                    admissible_min[key] = value

    def _create_condition(self, result: Dict) -> bool:
        """create thread condition"""
        if len(self._search_thread_pool) < 2:
            return True
        obj_median = np.median(
            [thread.obj_best1 for id, thread in self._search_thread_pool.items() if id]
        )
        return result[self._ls.metric] * self._ls.metric_op < obj_median

    def _clean(self, thread_id: int):
        """delete thread and increase admissible region if converged,
        merge local threads if they are close
        """
        assert thread_id
        todelete = set()
        for id in self._search_thread_pool:
            if id and id != thread_id:
                if self._inferior(id, thread_id):
                    todelete.add(id)
        for id in self._search_thread_pool:
            if id and id != thread_id:
                if self._inferior(thread_id, id):
                    todelete.add(thread_id)
                    break
        create_new = False
        if self._search_thread_pool[thread_id].converged:
            self._is_ls_ever_converged = True
            todelete.add(thread_id)
            self._expand_admissible_region(
                self._ls_bound_min,
                self._ls_bound_max,
                self._search_thread_pool[thread_id].space,
            )
            if self._candidate_start_points:
                if not self._started_from_given:
                    # remove start points whose perf is worse than the converged
                    obj = self._search_thread_pool[thread_id].obj_best1
                    worse = [
                        trial_id
                        for trial_id, r in self._candidate_start_points.items()
                        if r and r[self._ls.metric] * self._ls.metric_op >= obj
                    ]
                    # logger.info(f"remove candidate start points {worse} than {obj}")
                    for trial_id in worse:
                        del self._candidate_start_points[trial_id]
                if self._candidate_start_points and self._started_from_low_cost:
                    create_new = True
        for id in todelete:
            del self._search_thread_pool[id]
        if create_new:
            self._create_thread_from_best_candidate()

    def _create_thread_from_best_candidate(self):
        # find the best start point
        best_trial_id = None
        obj_best = None
        for trial_id, r in self._candidate_start_points.items():
            if r and (
                best_trial_id is None
                or r[self._ls.metric] * self._ls.metric_op < obj_best
            ):
                best_trial_id = trial_id
                obj_best = r[self._ls.metric] * self._ls.metric_op
        if best_trial_id:
            # create a new thread
            config = {}
            result = self._candidate_start_points[best_trial_id]
            for key, value in result.items():
                if key.startswith("config/"):
                    config[key[7:]] = value
            self._started_from_given = True
            del self._candidate_start_points[best_trial_id]
            self._create_thread(
                config, result, self._subspace.get(best_trial_id, self._ls.space)
            )

    def _expand_admissible_region(self, lower, upper, space):
        """expand the admissible region for the subspace `space`"""
        for key in upper:
            ub = upper[key]
            if isinstance(ub, list):
                choice = space[key]["_choice_"]
                self._expand_admissible_region(
                    lower[key][choice], upper[key][choice], space[key]
                )
            elif isinstance(ub, dict):
                self._expand_admissible_region(lower[key], ub, space[key])
            else:
                upper[key] += self._ls.STEPSIZE
                lower[key] -= self._ls.STEPSIZE

    def _inferior(self, id1: int, id2: int) -> bool:
        """whether thread id1 is inferior to id2"""
        t1 = self._search_thread_pool[id1]
        t2 = self._search_thread_pool[id2]
        if t1.obj_best1 < t2.obj_best2:
            return False
        elif t1.resource and t1.resource < t2.resource:
            return False
        elif t2.reach(t1):
            return True
        return False

    def on_trial_result(self, trial_id: str, result: Dict):
        """receive intermediate result."""
        if trial_id not in self._trial_proposed_by:
            return
        thread_id = self._trial_proposed_by[trial_id]
        if thread_id not in self._search_thread_pool:
            return
        if result and self._metric_constraints:
            result[self._metric + self.lagrange] = result[self._metric]
        self._search_thread_pool[thread_id].on_trial_result(trial_id, result)

    def suggest(self, trial_id: str) -> Optional[Dict]:
        """choose thread, suggest a valid config."""
        if self._init_used and not self._points_to_evaluate:
            choice, backup = self._select_thread()
            # if choice < 0:  # timeout
            #     return None
            config = self._search_thread_pool[choice].suggest(trial_id)
            if not choice and config is not None and self._ls.resource:
                config[self._ls.prune_attr] = self.best_resource
            elif choice and config is None:
                # local search thread finishes
                if self._search_thread_pool[choice].converged:
                    self._expand_admissible_region(
                        self._ls_bound_min,
                        self._ls_bound_max,
                        self._search_thread_pool[choice].space,
                    )
                    del self._search_thread_pool[choice]
                return None
            # preliminary check; not checking config validation
            space = self._search_thread_pool[choice].space
            skip = self._should_skip(choice, trial_id, config, space)
            use_rs = 0
            if skip:
                if choice:
                    return None
                # use rs when BO fails to suggest a config
                config, space = self._ls.complete_config({})
                skip = self._should_skip(-1, trial_id, config, space)
                if skip:
                    return None
                use_rs = 1
            if choice or self._valid(
                config,
                self._ls.space,
                space,
                self._gs_admissible_min,
                self._gs_admissible_max,
            ):
                # LS or valid or no backup choice
                self._trial_proposed_by[trial_id] = choice
                self._search_thread_pool[choice].running += use_rs
            else:  # invalid config proposed by GS
                if choice == backup:
                    # use CFO's init point
                    init_config = self._ls.init_config
                    config, space = self._ls.complete_config(
                        init_config, self._ls_bound_min, self._ls_bound_max
                    )
                    self._trial_proposed_by[trial_id] = choice
                    self._search_thread_pool[choice].running += 1
                else:
                    thread = self._search_thread_pool[backup]
                    config = thread.suggest(trial_id)
                    space = thread.space
                    skip = self._should_skip(backup, trial_id, config, space)
                    if skip:
                        return None
                    self._trial_proposed_by[trial_id] = backup
                    choice = backup
            if not choice:  # global search
                # temporarily relax admissible region for parallel proposals
                self._update_admissible_region(
                    config,
                    self._gs_admissible_min,
                    self._gs_admissible_max,
                    space,
                    self._ls.space,
                )
            else:
                self._update_admissible_region(
                    config,
                    self._ls_bound_min,
                    self._ls_bound_max,
                    space,
                    self._ls.space,
                )
                self._gs_admissible_min.update(self._ls_bound_min)
                self._gs_admissible_max.update(self._ls_bound_max)
            signature = self._ls.config_signature(config, space)
            self._result[signature] = {}
            self._subspace[trial_id] = space
        else:  # use init config
            if self._candidate_start_points is not None and self._points_to_evaluate:
                self._candidate_start_points[trial_id] = None
            reward = None
            if self._points_to_evaluate:
                init_config = self._points_to_evaluate.pop(0)
                if self._evaluated_rewards:
                    reward = self._evaluated_rewards.pop(0)
            else:
                init_config = self._ls.init_config
            config, space = self._ls.complete_config(
                init_config, self._ls_bound_min, self._ls_bound_max
            )
            if reward is None:
                config_signature = self._ls.config_signature(config, space)
                result = self._result.get(config_signature)
                if result:  # tried before
                    return None
                elif result is None:  # not tried before
                    self._result[config_signature] = {}
                else:  # running but no result yet
                    return None
            self._init_used = True
            self._trial_proposed_by[trial_id] = 0
            self._search_thread_pool[0].running += 1
            self._subspace[trial_id] = space
            if reward is not None:
                result = {self._metric: reward, self.cost_attr: 1, "config": config}
                self.on_trial_complete(trial_id, result)
                return None
        return config

    def _should_skip(self, choice, trial_id, config, space) -> bool:
        """if config is None or config's result is known or constraints are violated
        return True; o.w. return False
        """
        if config is None:
            return True
        config_signature = self._ls.config_signature(config, space)
        exists = config_signature in self._result
        # check constraints
        if not exists and self._config_constraints:
            for constraint in self._config_constraints:
                func, sign, threshold = constraint
                value = func(config)
                if (
                    sign == "<="
                    and value > threshold
                    or sign == ">="
                    and value < threshold
                ):
                    self._result[config_signature] = {
                        self._metric: np.inf * self._ls.metric_op,
                        "time_total_s": 1,
                    }
                    exists = True
                    break
        if exists:  # suggested before
            if choice >= 0:  # not fallback to rs
                result = self._result.get(config_signature)
                if result:  # finished
                    self._search_thread_pool[choice].on_trial_complete(
                        trial_id, result, error=False
                    )
                    if choice:
                        # local search thread
                        self._clean(choice)
                # else:     # running
                #     # tell the thread there is an error
                #     self._search_thread_pool[choice].on_trial_complete(
                #         trial_id, {}, error=True)
            return True
        return False

    def _select_thread(self) -> Tuple:
        """thread selector; use can_suggest to check LS availability"""
        # update priority
        now = time.time()
        min_eci = self._deadline - now
        if min_eci <= 0:
            # return -1, -1
            # keep proposing new configs assuming no budget left
            min_eci = 0
        elif self._num_samples and self._num_samples > 0:
            # estimate time left according to num_samples limitation
            num_finished = len(self._result)
            num_proposed = num_finished + len(self._trial_proposed_by)
            num_left = max(self._num_samples - num_proposed, 0)
            if num_proposed > 0:
                time_used = now - self._start_time + self._time_used
                min_eci = min(min_eci, time_used / num_finished * num_left)
            # print(f"{min_eci}, {time_used / num_finished * num_left}, {num_finished}, {num_left}")
        max_speed = 0
        for thread in self._search_thread_pool.values():
            if thread.speed > max_speed:
                max_speed = thread.speed
        for thread in self._search_thread_pool.values():
            thread.update_eci(self._metric_target, max_speed)
            if thread.eci < min_eci:
                min_eci = thread.eci
        for thread in self._search_thread_pool.values():
            thread.update_priority(min_eci)

        top_thread_id = backup_thread_id = 0
        priority1 = priority2 = self._search_thread_pool[0].priority
        for thread_id, thread in self._search_thread_pool.items():
            if thread_id and thread.can_suggest:
                priority = thread.priority
                if priority > priority1:
                    priority1 = priority
                    top_thread_id = thread_id
                if priority > priority2 or backup_thread_id == 0:
                    priority2 = priority
                    backup_thread_id = thread_id
        return top_thread_id, backup_thread_id

    def _valid(
        self, config: Dict, space: Dict, subspace: Dict, lower: Dict, upper: Dict
    ) -> bool:
        """config validator"""
        normalized_config = normalize(config, subspace, config, {})
        for key, lb in lower.items():
            if key in config:
                value = normalized_config[key]
                if isinstance(lb, list):
                    domain = space[key]
                    index = indexof(domain, value)
                    nestedspace = subspace[key]
                    lb = lb[index]
                    ub = upper[key][index]
                elif isinstance(lb, dict):
                    nestedspace = subspace[key]
                    domain = space[key]
                    ub = upper[key]
                else:
                    nestedspace = None
                if nestedspace:
                    valid = self._valid(value, domain, nestedspace, lb, ub)
                    if not valid:
                        return False
                elif (
                    value + self._ls.STEPSIZE < lower[key]
                    or value > upper[key] + self._ls.STEPSIZE
                ):
                    return False
        return True


try:
    from ray import __version__ as ray_version

    assert ray_version >= "1.0.0"
    from ray.tune import (
        uniform,
        quniform,
        choice,
        randint,
        qrandint,
        randn,
        qrandn,
        loguniform,
        qloguniform,
    )
except (ImportError, AssertionError):
    from ..tune.sample import (
        uniform,
        quniform,
        choice,
        randint,
        qrandint,
        randn,
        qrandn,
        loguniform,
        qloguniform,
    )

try:
    from nni.tuner import Tuner as NNITuner
    from nni.utils import extract_scalar_reward
except ImportError:

    class NNITuner:
        pass

    def extract_scalar_reward(x: Dict):
        return x.get("default")


class BlendSearchTuner(BlendSearch, NNITuner):
    """Tuner class for NNI."""

    def receive_trial_result(self, parameter_id, parameters, value, **kwargs):
        """Receive trial's final result.

        Args:
            parameter_id: int.
            parameters: object created by `generate_parameters()`.
            value: final metrics of the trial, including default metric.
        """
        result = {
            "config": parameters,
            self._metric: extract_scalar_reward(value),
            self.cost_attr: 1
            if isinstance(value, float)
            else value.get(self.cost_attr, value.get("sequence", 1))
            # if nni does not report training cost,
            # using sequence as an approximation.
            # if no sequence, using a constant 1
        }
        self.on_trial_complete(str(parameter_id), result)

    ...

    def generate_parameters(self, parameter_id, **kwargs) -> Dict:
        """Returns a set of trial (hyper-)parameters, as a serializable object.

        Args:
            parameter_id: int.
        """
        return self.suggest(str(parameter_id))

    ...

    def update_search_space(self, search_space):
        """Required by NNI.

        Tuners are advised to support updating search space at run-time.
        If a tuner can only set search space once before generating first hyper-parameters,
        it should explicitly document this behaviour.

        Args:
            search_space: JSON object created by experiment owner.
        """
        config = {}
        for key, value in search_space.items():
            v = value.get("_value")
            _type = value["_type"]
            if _type == "choice":
                config[key] = choice(v)
            elif _type == "randint":
                config[key] = randint(*v)
            elif _type == "uniform":
                config[key] = uniform(*v)
            elif _type == "quniform":
                config[key] = quniform(*v)
            elif _type == "loguniform":
                config[key] = loguniform(*v)
            elif _type == "qloguniform":
                config[key] = qloguniform(*v)
            elif _type == "normal":
                config[key] = randn(*v)
            elif _type == "qnormal":
                config[key] = qrandn(*v)
            else:
                raise ValueError(f"unsupported type in search_space {_type}")
        # low_cost_partial_config is passed to constructor,
        # which is before update_search_space() is called
        init_config = self._ls.init_config
        add_cost_to_space(config, init_config, self._cat_hp_cost)
        self._ls = self.LocalSearch(
            init_config,
            self._ls.metric,
            self._mode,
            config,
            self._ls.prune_attr,
            self._ls.min_resource,
            self._ls.max_resource,
            self._ls.resource_multiple_factor,
            cost_attr=self.cost_attr,
            seed=self._ls.seed,
        )
        if self._gs is not None:
            self._gs = GlobalSearch(
                space=config,
                metric=self._metric,
                mode=self._mode,
                sampler=self._gs._sampler,
            )
            self._gs.space = config
        self._init_search()


class CFO(BlendSearchTuner):
    """class for CFO algorithm."""

    __name__ = "CFO"

    def suggest(self, trial_id: str) -> Optional[Dict]:
        # Number of threads is 1 or 2. Thread 0 is a vacuous thread
        assert len(self._search_thread_pool) < 3, len(self._search_thread_pool)
        if len(self._search_thread_pool) < 2:
            # When a local thread converges, the number of threads is 1
            # Need to restart
            self._init_used = False
        return super().suggest(trial_id)

    def _select_thread(self) -> Tuple:
        for key in self._search_thread_pool:
            if key:
                return key, key

    def _create_condition(self, result: Dict) -> bool:
        """create thread condition"""
        if self._points_to_evaluate:
            # still evaluating user-specified init points
            # we evaluate all candidate start points before we
            # create the first local search thread
            return False
        if len(self._search_thread_pool) == 2:
            return False
        if self._candidate_start_points and self._thread_count == 1:
            # result needs to match or exceed the best candidate start point
            obj_best = min(
                (
                    self._ls.metric_op * r[self._ls.metric]
                    for r in self._candidate_start_points.values()
                    if r
                ),
                default=-np.inf,
            )

            return result[self._ls.metric] * self._ls.metric_op <= obj_best
        else:
            return True

    def on_trial_complete(
        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
    ):
        super().on_trial_complete(trial_id, result, error)
        if self._candidate_start_points and trial_id in self._candidate_start_points:
            # the trial is a candidate start point
            self._candidate_start_points[trial_id] = result
            if len(self._search_thread_pool) < 2 and not self._points_to_evaluate:
                self._create_thread_from_best_candidate()


class RandomSearch(CFO):
    """Class for random search."""

    def suggest(self, trial_id: str) -> Optional[Dict]:
        if self._points_to_evaluate:
            return super().suggest(trial_id)
        config, _ = self._ls.complete_config({})
        return config

    def on_trial_complete(
        self, trial_id: str, result: Optional[Dict] = None, error: bool = False
    ):
        return

    def on_trial_result(self, trial_id: str, result: Dict):
        return