[tune] added type hints (#10806)

Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
This commit is contained in:
Kai Fricke
2020-09-15 21:03:56 -07:00
committed by GitHub
co-authored by Richard Liaw
parent 5e030db8a5
commit c9fafe7733
31 changed files with 709 additions and 472 deletions
+13 -7
View File
@@ -1,5 +1,8 @@
from typing import Dict, List, Optional
from ray.tune.suggest.suggestion import Searcher, ConcurrencyLimiter
from ray.tune.suggest.search_generator import SearchGenerator
from ray.tune.trial import Trial
class _MockSearcher(Searcher):
@@ -11,29 +14,32 @@ class _MockSearcher(Searcher):
self.results = []
super(_MockSearcher, self).__init__(**kwargs)
def suggest(self, trial_id):
def suggest(self, trial_id: str):
if not self.stall:
self.live_trials[trial_id] = 1
return {"test_variable": 2}
return None
def on_trial_result(self, trial_id, result):
def on_trial_result(self, trial_id: str, result: Dict):
self.counter["result"] += 1
self.results += [result]
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = 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):
def _process_result(self, result: Dict):
self.final_results += [result]
class _MockSuggestionAlgorithm(SearchGenerator):
def __init__(self, max_concurrent=None, **kwargs):
def __init__(self, max_concurrent: Optional[int] = None, **kwargs):
self.searcher = _MockSearcher(**kwargs)
if max_concurrent:
self.searcher = ConcurrencyLimiter(
@@ -41,9 +47,9 @@ class _MockSuggestionAlgorithm(SearchGenerator):
super(_MockSuggestionAlgorithm, self).__init__(self.searcher)
@property
def live_trials(self):
def live_trials(self) -> List[Trial]:
return self.searcher.live_trials
@property
def results(self):
def results(self) -> List[Dict]:
return self.searcher.results
+12 -11
View File
@@ -1,4 +1,4 @@
from typing import Dict
from typing import Dict, List, Optional
from ax.service.ax_client import AxClient
from ray.tune.sample import Categorical, Float, Integer, LogUniform, \
@@ -103,14 +103,14 @@ class AxSearch(Searcher):
"""
def __init__(self,
space=None,
metric=None,
mode=None,
parameter_constraints=None,
outcome_constraints=None,
ax_client=None,
use_early_stopped_trials=None,
max_concurrent=None):
space: Optional[List[Dict]] = None,
metric: Optional[str] = None,
mode: Optional[str] = None,
parameter_constraints: Optional[List] = None,
outcome_constraints: Optional[List] = None,
ax_client: Optional[AxClient] = None,
use_early_stopped_trials: Optional[bool] = None,
max_concurrent: Optional[int] = None):
assert ax is not None, "Ax must be installed!"
if mode:
assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."
@@ -177,7 +177,8 @@ class AxSearch(Searcher):
logger.warning("Detected sequential enforcement. Be sure to use "
"a ConcurrencyLimiter.")
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict):
if self._ax:
return False
space = self.convert_search_space(config)
@@ -189,7 +190,7 @@ class AxSearch(Searcher):
self.setup_experiment()
return True
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
if not self._ax:
raise RuntimeError(
"Trying to sample a configuration from {}, but no search "
+6 -3
View File
@@ -2,9 +2,10 @@ import itertools
import os
import random
import uuid
from typing import Dict, List, Union
from ray.tune.error import TuneError
from ray.tune.experiment import convert_to_experiment_list
from ray.tune.experiment import Experiment, convert_to_experiment_list
from ray.tune.config_parser import make_parser, create_trial_from_spec
from ray.tune.suggest.variant_generator import (generate_variants, format_vars,
flatten_resolved_vars)
@@ -42,7 +43,7 @@ class BasicVariantGenerator(SearchAlgorithm):
searcher.is_finished == True
"""
def __init__(self, shuffle=False):
def __init__(self, shuffle: bool = False):
"""Initializes the Variant Generator.
"""
@@ -60,7 +61,9 @@ class BasicVariantGenerator(SearchAlgorithm):
else:
self._uuid_prefix = str(uuid.uuid1().hex)[:5] + "_"
def add_configurations(self, experiments):
def add_configurations(
self,
experiments: Union[Experiment, List[Experiment], Dict[str, Dict]]):
"""Chains generator given experiment specifications.
Arguments:
+28 -23
View File
@@ -2,9 +2,10 @@ from collections import defaultdict
import logging
import pickle
import json
from typing import Dict
from typing import Dict, Optional, Tuple
from ray.tune.sample import Float, Quantized
from ray.tune import ExperimentAnalysis
from ray.tune.sample import Domain, Float, Quantized
from ray.tune.suggest.variant_generator import parse_spec_vars
from ray.tune.utils.util import unflatten_dict
@@ -100,18 +101,18 @@ class BayesOptSearch(Searcher):
optimizer = None
def __init__(self,
space=None,
metric=None,
mode=None,
utility_kwargs=None,
random_state=42,
random_search_steps=10,
verbose=0,
patience=5,
skip_duplicate=True,
analysis=None,
max_concurrent=None,
use_early_stopped_trials=None):
space: Optional[Dict] = None,
metric: Optional[str] = None,
mode: Optional[str] = None,
utility_kwargs: Optional[Dict] = None,
random_state: int = 42,
random_search_steps: int = 10,
verbose: int = 0,
patience: int = 5,
skip_duplicate: bool = True,
analysis: Optional[ExperimentAnalysis] = None,
max_concurrent: Optional[int] = None,
use_early_stopped_trials: Optional[bool] = None):
"""Instantiate new BayesOptSearch object.
Args:
@@ -200,7 +201,8 @@ class BayesOptSearch(Searcher):
verbose=self._verbose,
random_state=self._random_state)
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
if self.optimizer:
return False
space = self.convert_search_space(config)
@@ -218,7 +220,7 @@ class BayesOptSearch(Searcher):
self.setup_optimizer()
return True
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
"""Return new point to be explored by black box function.
Args:
@@ -278,7 +280,7 @@ class BayesOptSearch(Searcher):
# Return a deep copy of the mapping
return unflatten_dict(config)
def register_analysis(self, analysis):
def register_analysis(self, analysis: ExperimentAnalysis):
"""Integrate the given analysis into the gaussian process.
Args:
@@ -293,7 +295,10 @@ class BayesOptSearch(Searcher):
# gaussian process optimizer
self._register_result(params, report)
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
"""Notification for the completion of trial.
Args:
@@ -330,18 +335,18 @@ class BayesOptSearch(Searcher):
for params, result in self._buffered_trial_results:
self._register_result(params, result)
def _register_result(self, params, result):
def _register_result(self, params: Tuple[str], result: Dict):
"""Register given tuple of params and results."""
self.optimizer.register(params, self._metric_op * result[self.metric])
def save(self, checkpoint_path):
def save(self, checkpoint_path: str):
"""Storing current optimizer state."""
with open(checkpoint_path, "wb") as f:
pickle.dump(
(self.optimizer, self._buffered_trial_results,
self._total_random_search_trials, self._config_counter), f)
def restore(self, checkpoint_path):
def restore(self, checkpoint_path: str):
"""Restoring current optimizer state."""
with open(checkpoint_path, "rb") as f:
(self.optimizer, self._buffered_trial_results,
@@ -349,7 +354,7 @@ class BayesOptSearch(Searcher):
self._config_counter) = pickle.load(f)
@staticmethod
def convert_search_space(spec: Dict):
def convert_search_space(spec: Dict) -> Dict:
spec = flatten_dict(spec, prevent_delimiter=True)
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
@@ -358,7 +363,7 @@ class BayesOptSearch(Searcher):
"Grid search parameters cannot be automatically converted "
"to a BayesOpt search space.")
def resolve_value(domain):
def resolve_value(domain: Domain) -> Tuple[float, float]:
sampler = domain.get_sampler()
if isinstance(sampler, Quantized):
logger.warning(
+23 -17
View File
@@ -3,10 +3,11 @@
import copy
import logging
import math
from typing import Dict
from typing import Dict, Optional
import ConfigSpace
from ray.tune.sample import Categorical, Float, Integer, LogUniform, Normal, \
from ray.tune.sample import Categorical, Domain, Float, Integer, LogUniform, \
Normal, \
Quantized, \
Uniform
from ray.tune.suggest import Searcher
@@ -20,7 +21,7 @@ logger = logging.getLogger(__name__)
class _BOHBJobWrapper():
"""Mock object for HpBandSter to process."""
def __init__(self, loss, budget, config):
def __init__(self, loss: float, budget: float, config: Dict):
self.result = {"loss": loss}
self.kwargs = {"budget": budget, "config": config.copy()}
self.exception = None
@@ -92,11 +93,11 @@ class TuneBOHB(Searcher):
"""
def __init__(self,
space=None,
bohb_config=None,
max_concurrent=10,
metric=None,
mode=None):
space: Optional[ConfigSpace.ConfigurationSpace] = None,
bohb_config: Optional[Dict] = None,
max_concurrent: int = 10,
metric: Optional[str] = None,
mode: Optional[str] = None):
from hpbandster.optimizers.config_generators.bohb import BOHB
assert BOHB is not None, "HpBandSter must be installed!"
if mode:
@@ -126,7 +127,8 @@ class TuneBOHB(Searcher):
bohb_config = self._bohb_config or {}
self.bohber = BOHB(self._space, **bohb_config)
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
if self._space:
return False
space = self.convert_search_space(config)
@@ -140,7 +142,7 @@ class TuneBOHB(Searcher):
self.setup_bohb()
return True
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
if not self._space:
raise RuntimeError(
"Trying to sample a configuration from {}, but no search "
@@ -156,7 +158,7 @@ class TuneBOHB(Searcher):
return unflatten_dict(config)
return None
def on_trial_result(self, trial_id, result):
def on_trial_result(self, trial_id: str, result: Dict):
if trial_id not in self.paused:
self.running.add(trial_id)
if "hyperband_info" not in result:
@@ -166,28 +168,31 @@ class TuneBOHB(Searcher):
hbs_wrapper = self.to_wrapper(trial_id, result)
self.bohber.new_result(hbs_wrapper)
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
del self.trial_to_params[trial_id]
if trial_id in self.paused:
self.paused.remove(trial_id)
if trial_id in self.running:
self.running.remove(trial_id)
def to_wrapper(self, trial_id, result):
def to_wrapper(self, trial_id: str, result: Dict) -> _BOHBJobWrapper:
return _BOHBJobWrapper(self._metric_op * result[self.metric],
result["hyperband_info"]["budget"],
self.trial_to_params[trial_id])
def on_pause(self, trial_id):
def on_pause(self, trial_id: str):
self.paused.add(trial_id)
self.running.remove(trial_id)
def on_unpause(self, trial_id):
def on_unpause(self, trial_id: str):
self.paused.remove(trial_id)
self.running.add(trial_id)
@staticmethod
def convert_search_space(spec: Dict):
def convert_search_space(spec: Dict) -> ConfigSpace.ConfigurationSpace:
spec = flatten_dict(spec, prevent_delimiter=True)
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
@@ -196,7 +201,8 @@ class TuneBOHB(Searcher):
"Grid search parameters cannot be automatically converted "
"to a TuneBOHB search space.")
def resolve_value(par, domain):
def resolve_value(par: str, domain: Domain
) -> ConfigSpace.hyperparameters.Hyperparameter:
quantize = None
sampler = domain.get_sampler()
+23 -19
View File
@@ -5,16 +5,18 @@ from __future__ import print_function
import inspect
import logging
import pickle
from typing import Dict
from typing import Dict, List, Optional
from ray.tune.sample import Float, Quantized
from ray.tune.sample import Domain, Float, Quantized
from ray.tune.suggest.variant_generator import parse_spec_vars
from ray.tune.utils.util import flatten_dict
try: # Python 3 only -- needed for lint test.
import dragonfly
from dragonfly.opt.blackbox_optimiser import BlackboxOptimiser
except ImportError:
dragonfly = None
BlackboxOptimiser = None
from ray.tune.suggest.suggestion import Searcher
@@ -127,13 +129,13 @@ class DragonflySearch(Searcher):
"""
def __init__(self,
optimizer=None,
domain=None,
space=None,
metric=None,
mode=None,
points_to_evaluate=None,
evaluated_rewards=None,
optimizer: Optional[BlackboxOptimiser] = None,
domain: Optional[str] = None,
space: Optional[List[Dict]] = None,
metric: Optional[str] = None,
mode: Optional[str] = None,
points_to_evaluate: Optional[List[List]] = None,
evaluated_rewards: Optional[List] = None,
**kwargs):
assert dragonfly is not None, """dragonfly must be installed!
You can install Dragonfly with the command:
@@ -144,8 +146,6 @@ class DragonflySearch(Searcher):
super(DragonflySearch, self).__init__(
metric=metric, mode=mode, **kwargs)
from dragonfly.opt.blackbox_optimiser import BlackboxOptimiser
self._opt_arg = optimizer
self._domain = domain
self._space = space
@@ -245,7 +245,8 @@ class DragonflySearch(Searcher):
elif self._mode == "max":
self._metric_op = 1.
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
if self._opt:
return False
space = self.convert_search_space(config)
@@ -258,7 +259,7 @@ class DragonflySearch(Searcher):
self.setup_dragonfly()
return True
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
if not self._opt:
raise RuntimeError(
"Trying to sample a configuration from {}, but no search "
@@ -281,26 +282,29 @@ class DragonflySearch(Searcher):
self._live_trial_mapping[trial_id] = suggested_config
return {"point": suggested_config}
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
"""Passes result to Dragonfly unless early terminated or errored."""
trial_info = self._live_trial_mapping.pop(trial_id)
if result:
self._opt.tell([(trial_info,
self._metric_op * result[self._metric])])
def save(self, checkpoint_dir):
def save(self, checkpoint_path: str):
trials_object = (self._initial_points, self._opt)
with open(checkpoint_dir, "wb") as outputFile:
with open(checkpoint_path, "wb") as outputFile:
pickle.dump(trials_object, outputFile)
def restore(self, checkpoint_dir):
def restore(self, checkpoint_dir: str):
with open(checkpoint_dir, "rb") as inputFile:
trials_object = pickle.load(inputFile)
self._initial_points = trials_object[0]
self._opt = trials_object[1]
@staticmethod
def convert_search_space(spec: Dict):
def convert_search_space(spec: Dict) -> List[Dict]:
spec = flatten_dict(spec, prevent_delimiter=True)
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
@@ -309,7 +313,7 @@ class DragonflySearch(Searcher):
"Grid search parameters cannot be automatically converted "
"to a Dragonfly search space.")
def resolve_value(par, domain):
def resolve_value(par: str, domain: Domain) -> Dict:
sampler = domain.get_sampler()
if isinstance(sampler, Quantized):
logger.warning(
+30 -25
View File
@@ -1,4 +1,4 @@
from typing import Dict
from typing import Any, Dict, List, Optional
import numpy as np
import copy
@@ -6,7 +6,8 @@ import logging
from functools import partial
import pickle
from ray.tune.sample import Categorical, Float, Integer, LogUniform, Normal, \
from ray.tune.sample import Categorical, Domain, Float, Integer, LogUniform, \
Normal, \
Quantized, \
Uniform
from ray.tune.suggest.variant_generator import assign_value, parse_spec_vars
@@ -117,15 +118,15 @@ class HyperOptSearch(Searcher):
def __init__(
self,
space=None,
metric=None,
mode=None,
points_to_evaluate=None,
n_initial_points=20,
random_state_seed=None,
gamma=0.25,
max_concurrent=None,
use_early_stopped_trials=None,
space: Optional[Dict] = None,
metric: Optional[str] = None,
mode: Optional[str] = None,
points_to_evaluate: Optional[List[Dict]] = None,
n_initial_points: int = 20,
random_state_seed: Optional[int] = None,
gamma: float = 0.25,
max_concurrent: Optional[int] = None,
use_early_stopped_trials: Optional[bool] = None,
):
assert hpo is not None, (
"HyperOpt must be installed! Run `pip install hyperopt`.")
@@ -170,7 +171,8 @@ class HyperOptSearch(Searcher):
if space:
self.domain = hpo.Domain(lambda spc: spc, space)
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
if self.domain:
return False
space = self.convert_search_space(config)
@@ -188,7 +190,7 @@ class HyperOptSearch(Searcher):
return True
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
if not self.domain:
raise RuntimeError(
"Trying to sample a configuration from {}, but no search "
@@ -226,7 +228,7 @@ class HyperOptSearch(Searcher):
print_node_on_error=self.domain.rec_eval_print_node_on_error)
return copy.deepcopy(suggested_config)
def on_trial_result(self, trial_id, result):
def on_trial_result(self, trial_id: str, result: Dict):
ho_trial = self._get_hyperopt_trial(trial_id)
if ho_trial is None:
return
@@ -234,7 +236,10 @@ class HyperOptSearch(Searcher):
ho_trial["book_time"] = now
ho_trial["refresh_time"] = now
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
"""Notification for the completion of trial.
The result is internally negated when interacting with HyperOpt
@@ -252,7 +257,7 @@ class HyperOptSearch(Searcher):
self._process_result(trial_id, result)
del self._live_trial_mapping[trial_id]
def _process_result(self, trial_id, result):
def _process_result(self, trial_id: str, result: Dict):
ho_trial = self._get_hyperopt_trial(trial_id)
if not ho_trial:
return
@@ -263,10 +268,10 @@ class HyperOptSearch(Searcher):
ho_trial["result"] = hp_result
self._hpopt_trials.refresh()
def _to_hyperopt_result(self, result):
def _to_hyperopt_result(self, result: Dict) -> Dict:
return {"loss": self.metric_op * result[self.metric], "status": "ok"}
def _get_hyperopt_trial(self, trial_id):
def _get_hyperopt_trial(self, trial_id: str) -> Optional[Dict]:
if trial_id not in self._live_trial_mapping:
return
hyperopt_tid = self._live_trial_mapping[trial_id][0]
@@ -274,21 +279,21 @@ class HyperOptSearch(Searcher):
t for t in self._hpopt_trials.trials if t["tid"] == hyperopt_tid
][0]
def get_state(self):
def get_state(self) -> Dict:
return {
"hyperopt_trials": self._hpopt_trials,
"rstate": self.rstate.get_state()
}
def set_state(self, state):
def set_state(self, state: Dict):
self._hpopt_trials = state["hyperopt_trials"]
self.rstate.set_state(state["rstate"])
def save(self, checkpoint_path):
def save(self, checkpoint_path: str):
with open(checkpoint_path, "wb") as outputFile:
pickle.dump(self.get_state(), outputFile)
def restore(self, checkpoint_path):
def restore(self, checkpoint_path: str):
with open(checkpoint_path, "rb") as inputFile:
trials_object = pickle.load(inputFile)
@@ -299,19 +304,19 @@ class HyperOptSearch(Searcher):
self.set_state(trials_object)
@staticmethod
def convert_search_space(spec: Dict):
def convert_search_space(spec: Dict) -> Dict:
spec = copy.deepcopy(spec)
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
if not domain_vars and not grid_vars:
return []
return {}
if grid_vars:
raise ValueError(
"Grid search parameters cannot be automatically converted "
"to a HyperOpt search space.")
def resolve_value(par, domain):
def resolve_value(par: str, domain: Domain) -> Any:
quantize = None
sampler = domain.get_sampler()
+28 -17
View File
@@ -1,16 +1,23 @@
import logging
import pickle
from typing import Dict
from typing import Dict, Optional, Union
from ray.tune.sample import Categorical, Float, Integer, LogUniform, Quantized
from ray.tune.sample import Categorical, Domain, Float, Integer, LogUniform, \
Quantized
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 nevergrad as ng
from nevergrad.optimization import Optimizer
from nevergrad.optimization.base import ConfiguredOptimizer
Parameter = ng.p.Parameter
except ImportError:
ng = None
Optimizer = None
ConfiguredOptimizer = None
Parameter = None
from ray.tune.suggest import Searcher
@@ -85,11 +92,11 @@ class NevergradSearch(Searcher):
"""
def __init__(self,
optimizer=None,
space=None,
metric=None,
mode=None,
max_concurrent=None,
optimizer: Union[None, Optimizer, ConfiguredOptimizer] = None,
space: Optional[Parameter] = None,
metric: Optional[str] = None,
mode: Optional[str] = None,
max_concurrent: Optional[int] = None,
**kwargs):
assert ng is not None, "Nevergrad must be installed!"
if mode:
@@ -102,7 +109,7 @@ class NevergradSearch(Searcher):
self._opt_factory = None
self._nevergrad_opt = None
if isinstance(optimizer, ng.optimization.Optimizer):
if isinstance(optimizer, Optimizer):
if space is not None or isinstance(space, list):
raise ValueError(
"If you pass a configured optimizer to Nevergrad, either "
@@ -110,7 +117,7 @@ class NevergradSearch(Searcher):
"parameter.")
self._parameters = space
self._nevergrad_opt = optimizer
elif isinstance(optimizer, ng.optimization.base.ConfiguredOptimizer):
elif isinstance(optimizer, ConfiguredOptimizer):
self._opt_factory = optimizer
self._parameters = None
self._space = space
@@ -155,7 +162,8 @@ class NevergradSearch(Searcher):
raise ValueError("len(parameters_names) must match optimizer "
"dimension for non-instrumented optimizers")
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
if self._nevergrad_opt or self._space:
return False
space = self.convert_search_space(config)
@@ -169,7 +177,7 @@ class NevergradSearch(Searcher):
self.setup_nevergrad()
return True
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
if not self._nevergrad_opt:
raise RuntimeError(
"Trying to sample a configuration from {}, but no search "
@@ -192,7 +200,10 @@ class NevergradSearch(Searcher):
else:
return unflatten_dict(suggested_config.kwargs)
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
"""Notification for the completion of trial.
The result is internally negated when interacting with Nevergrad
@@ -204,24 +215,24 @@ class NevergradSearch(Searcher):
self._live_trial_mapping.pop(trial_id)
def _process_result(self, trial_id, result):
def _process_result(self, trial_id: str, result: Dict):
ng_trial_info = self._live_trial_mapping[trial_id]
self._nevergrad_opt.tell(ng_trial_info,
self._metric_op * result[self._metric])
def save(self, checkpoint_path):
def save(self, checkpoint_path: str):
trials_object = (self._nevergrad_opt, self._parameters)
with open(checkpoint_path, "wb") as outputFile:
pickle.dump(trials_object, outputFile)
def restore(self, checkpoint_path):
def restore(self, checkpoint_path: str):
with open(checkpoint_path, "rb") as inputFile:
trials_object = pickle.load(inputFile)
self._nevergrad_opt = trials_object[0]
self._parameters = trials_object[1]
@staticmethod
def convert_search_space(spec: Dict):
def convert_search_space(spec: Dict) -> Parameter:
spec = flatten_dict(spec, prevent_delimiter=True)
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
@@ -230,7 +241,7 @@ class NevergradSearch(Searcher):
"Grid search parameters cannot be automatically converted "
"to a Nevergrad search space.")
def resolve_value(domain):
def resolve_value(domain: Domain) -> Parameter:
sampler = domain.get_sampler()
if isinstance(sampler, Quantized):
logger.warning("Nevergrad does not support quantization. "
+23 -13
View File
@@ -1,9 +1,9 @@
import logging
import pickle
from typing import Dict
from typing import Dict, List, Optional, Tuple
from ray.tune.result import TRAINING_ITERATION
from ray.tune.sample import Categorical, Float, Integer, LogUniform, \
from ray.tune.sample import Categorical, Domain, Float, Integer, LogUniform, \
Quantized, Uniform
from ray.tune.suggest.variant_generator import parse_spec_vars
from ray.tune.utils import flatten_dict
@@ -11,8 +11,10 @@ from ray.tune.utils.util import unflatten_dict
try:
import optuna as ot
from optuna.samplers import BaseSampler
except ImportError:
ot = None
BaseSampler = None
from ray.tune.suggest import Searcher
@@ -100,7 +102,11 @@ class OptunaSearch(Searcher):
"""
def __init__(self, space=None, metric=None, mode=None, sampler=None):
def __init__(self,
space: Optional[List[Tuple]] = None,
metric: Optional[str] = None,
mode: Optional[str] = None,
sampler: Optional[BaseSampler] = None):
assert ot is not None, (
"Optuna must be installed! Run `pip install optuna`.")
super(OptunaSearch, self).__init__(
@@ -113,7 +119,7 @@ class OptunaSearch(Searcher):
self._study_name = "optuna" # Fixed study name for in-memory storage
self._sampler = sampler or ot.samplers.TPESampler()
assert isinstance(self._sampler, ot.samplers.BaseSampler), \
assert isinstance(self._sampler, BaseSampler), \
"You can only pass an instance of `optuna.samplers.BaseSampler` " \
"as a sampler to `OptunaSearcher`."
@@ -125,7 +131,7 @@ class OptunaSearch(Searcher):
if self._space:
self.setup_study(mode)
def setup_study(self, mode):
def setup_study(self, mode: str):
self._ot_study = ot.study.create_study(
storage=self._storage,
sampler=self._sampler,
@@ -134,7 +140,8 @@ class OptunaSearch(Searcher):
direction="minimize" if mode == "min" else "maximize",
load_if_exists=True)
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
if self._space:
return False
space = self.convert_search_space(config)
@@ -146,7 +153,7 @@ class OptunaSearch(Searcher):
self.setup_study(mode)
return True
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
if not self._space:
raise RuntimeError(
"Trying to sample a configuration from {}, but no search "
@@ -169,13 +176,16 @@ class OptunaSearch(Searcher):
}
return unflatten_dict(params)
def on_trial_result(self, trial_id, result):
def on_trial_result(self, trial_id: str, result: Dict):
metric = result[self.metric]
step = result[TRAINING_ITERATION]
ot_trial = self._ot_trials[trial_id]
ot_trial.report(metric, step)
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
ot_trial = self._ot_trials[trial_id]
ot_trial_id = ot_trial._trial_id
self._storage.set_trial_value(ot_trial_id, result.get(
@@ -183,20 +193,20 @@ class OptunaSearch(Searcher):
self._storage.set_trial_state(ot_trial_id,
ot.trial.TrialState.COMPLETE)
def save(self, checkpoint_path):
def save(self, checkpoint_path: str):
save_object = (self._storage, self._pruner, self._sampler,
self._ot_trials, self._ot_study)
with open(checkpoint_path, "wb") as outputFile:
pickle.dump(save_object, outputFile)
def restore(self, checkpoint_path):
def restore(self, checkpoint_path: str):
with open(checkpoint_path, "rb") as inputFile:
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):
def convert_search_space(spec: Dict) -> List[Tuple]:
spec = flatten_dict(spec, prevent_delimiter=True)
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
@@ -208,7 +218,7 @@ class OptunaSearch(Searcher):
"Grid search parameters cannot be automatically converted "
"to an Optuna search space.")
def resolve_value(par, domain):
def resolve_value(par: str, domain: Domain) -> Tuple:
quantize = None
sampler = domain.get_sampler()
+26 -14
View File
@@ -1,5 +1,7 @@
import copy
import logging
from typing import Dict, List, Optional
import numpy as np
from ray.tune.suggest.suggestion import Searcher
@@ -10,7 +12,7 @@ TRIAL_INDEX = "__trial_index__"
"""str: A constant value representing the repeat index of the trial."""
def _warn_num_samples(searcher, num_samples):
def _warn_num_samples(searcher: Searcher, num_samples: int):
if isinstance(searcher, Repeater) and num_samples % searcher.repeat:
logger.warning(
"`num_samples` is now expected to be the total number of trials, "
@@ -34,7 +36,10 @@ class _TrialGroup:
"""
def __init__(self, primary_trial_id, config, max_trials=1):
def __init__(self,
primary_trial_id: str,
config: Dict,
max_trials: int = 1):
assert type(config) is dict, (
"config is not a dict, got {}".format(config))
self.primary_trial_id = primary_trial_id
@@ -42,27 +47,27 @@ class _TrialGroup:
self._trials = {primary_trial_id: None}
self.max_trials = max_trials
def add(self, trial_id):
def add(self, trial_id: str):
assert len(self._trials) < self.max_trials
self._trials.setdefault(trial_id, None)
def full(self):
def full(self) -> bool:
return len(self._trials) == self.max_trials
def report(self, trial_id, score):
def report(self, trial_id: str, score: float):
assert trial_id in self._trials
if score is None:
raise ValueError("Internal Error: Score cannot be None.")
self._trials[trial_id] = score
def finished_reporting(self):
def finished_reporting(self) -> bool:
return None not in self._trials.values() and len(
self._trials) == self.max_trials
def scores(self):
def scores(self) -> List[Optional[float]]:
return list(self._trials.values())
def count(self):
def count(self) -> int:
return len(self._trials)
@@ -103,7 +108,10 @@ class Repeater(Searcher):
"""
def __init__(self, searcher, repeat=1, set_index=True):
def __init__(self,
searcher: Searcher,
repeat: int = 1,
set_index: bool = True):
self.searcher = searcher
self.repeat = repeat
self._set_index = set_index
@@ -113,7 +121,7 @@ class Repeater(Searcher):
super(Repeater, self).__init__(
metric=self.searcher.metric, mode=self.searcher.mode)
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
if self._current_group is None or self._current_group.full():
config = self.searcher.suggest(trial_id)
if config is None:
@@ -132,7 +140,10 @@ class Repeater(Searcher):
self._trial_id_to_group[trial_id] = self._current_group
return config
def on_trial_complete(self, trial_id, result=None, **kwargs):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
**kwargs):
"""Stores the score for and keeps track of a completed trial.
Stores the metric of a trial as nan if any of the following conditions
@@ -160,13 +171,14 @@ class Repeater(Searcher):
result={self.searcher.metric: np.nanmean(scores)},
**kwargs)
def get_state(self):
def get_state(self) -> Dict:
self_state = self.__dict__.copy()
del self_state["searcher"]
return self_state
def set_state(self, state):
def set_state(self, state: Dict):
self.__dict__.update(state)
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
return self.searcher.set_search_properties(metric, mode, config)
+22 -9
View File
@@ -1,3 +1,9 @@
from typing import Dict, List, Optional, Union
from ray.tune.experiment import Experiment
from ray.tune.trial import Trial
class SearchAlgorithm:
"""Interface of an event handler API for hyperparameter search.
@@ -12,7 +18,8 @@ class SearchAlgorithm:
"""
_finished = False
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
"""Pass search properties to search algorithm.
This method acts as an alternative to instantiating search algorithms
@@ -29,7 +36,9 @@ class SearchAlgorithm:
"""
return True
def add_configurations(self, experiments):
def add_configurations(
self,
experiments: Union[Experiment, List[Experiment], Dict[str, Dict]]):
"""Tracks given experiment specifications.
Arguments:
@@ -37,7 +46,7 @@ class SearchAlgorithm:
"""
raise NotImplementedError
def next_trials(self):
def next_trials(self) -> List[Trial]:
"""Provides Trial objects to be queued into the TrialRunner.
Returns:
@@ -45,17 +54,21 @@ class SearchAlgorithm:
"""
raise NotImplementedError
def on_trial_result(self, trial_id, result):
def on_trial_result(self, trial_id: str, result: Dict):
"""Called on each intermediate result returned by a trial.
This will only be called when the trial is in the RUNNING state.
Arguments:
trial_id: Identifier for the trial.
result: Result dictionary.
"""
pass
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
"""Notification for the completion of trial.
Arguments:
@@ -69,7 +82,7 @@ class SearchAlgorithm:
"""
pass
def is_finished(self):
def is_finished(self) -> bool:
"""Returns True if no trials left to be queued into TrialRunner.
Can return True before all trials have finished executing.
@@ -80,14 +93,14 @@ class SearchAlgorithm:
"""Marks the search algorithm as finished."""
self._finished = True
def has_checkpoint(self, dirpath):
def has_checkpoint(self, dirpath: str) -> bool:
"""Should return False if not restoring is not implemented."""
return False
def save_to_dir(self, dirpath, **kwargs):
def save_to_dir(self, dirpath: str, **kwargs):
"""Saves a search algorithm."""
pass
def restore_from_dir(self, dirpath):
def restore_from_dir(self, dirpath: str):
"""Restores a search algorithm along with its wrapped state."""
pass
+24 -16
View File
@@ -2,10 +2,11 @@ import os
import copy
import logging
import glob
from typing import Dict, List, Optional, Union
import ray.cloudpickle as cloudpickle
from ray.tune.error import TuneError
from ray.tune.experiment import convert_to_experiment_list
from ray.tune.experiment import Experiment, 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
@@ -21,7 +22,7 @@ def _warn_on_repeater(searcher, total_samples):
_warn_num_samples(searcher, total_samples)
def _atomic_save(state, checkpoint_dir, file_name):
def _atomic_save(state: Dict, checkpoint_dir: str, file_name: str):
"""Atomically saves the object to the checkpoint directory
This is automatically used by tune.run during a Tune job.
@@ -34,7 +35,7 @@ def _atomic_save(state, checkpoint_dir, file_name):
os.rename(tmp_search_ckpt_path, os.path.join(checkpoint_dir, file_name))
def _find_newest_ckpt(dirpath, pattern):
def _find_newest_ckpt(dirpath: str, pattern: str):
"""Returns path to most recently modified checkpoint."""
full_paths = glob.glob(os.path.join(dirpath, pattern))
if not full_paths:
@@ -58,7 +59,7 @@ class SearchGenerator(SearchAlgorithm):
"""
CKPT_FILE_TMPL = "search_gen_state-{}.json"
def __init__(self, searcher):
def __init__(self, searcher: Searcher):
assert issubclass(
type(searcher),
Searcher), ("Searcher should be subclassing Searcher.")
@@ -69,10 +70,13 @@ class SearchGenerator(SearchAlgorithm):
self._total_samples = None # int: total samples to evaluate.
self._finished = False
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
return self.searcher.set_search_properties(metric, mode, config)
def add_configurations(self, experiments):
def add_configurations(
self,
experiments: Union[Experiment, List[Experiment], Dict[str, Dict]]):
"""Registers experiment specifications.
Arguments:
@@ -91,7 +95,7 @@ class SearchGenerator(SearchAlgorithm):
if "run" not in experiment_spec:
raise TuneError("Must specify `run` in {}".format(experiment_spec))
def next_trials(self):
def next_trials(self) -> List[Trial]:
"""Provides a batch of Trial objects to be queued into the TrialRunner.
Returns:
@@ -106,7 +110,8 @@ class SearchGenerator(SearchAlgorithm):
trials.append(trial)
return trials
def create_trial_if_possible(self, experiment_spec, output_path):
def create_trial_if_possible(self, experiment_spec: Dict,
output_path: str) -> Optional[Trial]:
logger.debug("creating trial")
trial_id = Trial.generate_id()
suggested_config = self.searcher.suggest(trial_id)
@@ -135,18 +140,21 @@ class SearchGenerator(SearchAlgorithm):
trial_id=trial_id)
return trial
def on_trial_result(self, trial_id, result):
def on_trial_result(self, trial_id: str, result: Dict):
"""Notifies the underlying searcher."""
self.searcher.on_trial_result(trial_id, result)
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
self.searcher.on_trial_complete(
trial_id=trial_id, result=result, error=error)
def is_finished(self):
def is_finished(self) -> bool:
return self._counter >= self._total_samples or self._finished
def get_state(self):
def get_state(self) -> Dict:
return {
"counter": self._counter,
"total_samples": self._total_samples,
@@ -154,17 +162,17 @@ class SearchGenerator(SearchAlgorithm):
"experiment": self._experiment
}
def set_state(self, state):
def set_state(self, state: Dict):
self._counter = state["counter"]
self._total_samples = state["total_samples"]
self._finished = state["finished"]
self._experiment = state["experiment"]
def has_checkpoint(self, dirpath):
def has_checkpoint(self, dirpath: str):
return bool(
_find_newest_ckpt(dirpath, self.CKPT_FILE_TMPL.format("*")))
def save_to_dir(self, dirpath, session_str):
def save_to_dir(self, dirpath: str, session_str: str):
"""Saves self + searcher to dir.
Separates the "searcher" from its wrappers (concurrency, repeating).
@@ -196,7 +204,7 @@ class SearchGenerator(SearchAlgorithm):
_atomic_save(search_alg_state, dirpath,
self.CKPT_FILE_TMPL.format(session_str))
def restore_from_dir(self, dirpath):
def restore_from_dir(self, dirpath: str):
"""Restores self + searcher + search wrappers from dirpath."""
searcher = self.searcher
+23 -16
View File
@@ -2,10 +2,14 @@ import copy
import os
import logging
import pickle
from typing import Dict, List, Optional, Union
try:
import sigopt as sgo
Connection = sgo.Connection
except ImportError:
sgo = None
Connection = None
from ray.tune.suggest import Searcher
@@ -122,16 +126,16 @@ class SigOptSearch(Searcher):
}
def __init__(self,
space=None,
name="Default Tune Experiment",
max_concurrent=1,
reward_attr=None,
connection=None,
experiment_id=None,
observation_budget=None,
project=None,
metric="episode_reward_mean",
mode="max",
space: List[Dict] = None,
name: str = "Default Tune Experiment",
max_concurrent: int = 1,
reward_attr: Optional[str] = None,
connection: Optional[Connection] = None,
experiment_id: Optional[str] = None,
observation_budget: Optional[int] = None,
project: Optional[str] = None,
metric: Union[None, str, List[str]] = "episode_reward_mean",
mode: Union[None, str, List[str]] = "max",
**kwargs):
assert (experiment_id is
None) ^ (space is None), "space xor experiment_id must be set"
@@ -178,7 +182,7 @@ class SigOptSearch(Searcher):
super(SigOptSearch, self).__init__(metric=metric, mode=mode, **kwargs)
def suggest(self, trial_id):
def suggest(self, trial_id: str):
if self._max_concurrent:
if len(self._live_trial_mapping) >= self._max_concurrent:
return None
@@ -190,7 +194,10 @@ class SigOptSearch(Searcher):
return copy.deepcopy(suggestion.assignments)
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
"""Notification for the completion of trial.
If a trial fails, it will be reported as a failed Observation, telling
@@ -214,7 +221,7 @@ class SigOptSearch(Searcher):
del self._live_trial_mapping[trial_id]
@staticmethod
def serialize_metric(metrics, modes):
def serialize_metric(metrics: List[str], modes: List[str]):
"""
Converts metrics to https://app.sigopt.com/docs/objects/metric
"""
@@ -224,7 +231,7 @@ class SigOptSearch(Searcher):
dict(name=metric, **SigOptSearch.OBJECTIVE_MAP[mode].copy()))
return serialized_metric
def serialize_result(self, result):
def serialize_result(self, result: Dict):
"""
Converts experiments results to
https://app.sigopt.com/docs/objects/metric_evaluation
@@ -244,12 +251,12 @@ class SigOptSearch(Searcher):
values.append(value)
return values
def save(self, checkpoint_path):
def save(self, checkpoint_path: str):
trials_object = (self.conn, self.experiment)
with open(checkpoint_path, "wb") as outputFile:
pickle.dump(trials_object, outputFile)
def restore(self, checkpoint_path):
def restore(self, checkpoint_path: str):
with open(checkpoint_path, "rb") as inputFile:
trials_object = pickle.load(inputFile)
self.conn = trials_object[0]
+26 -21
View File
@@ -1,8 +1,8 @@
import logging
import pickle
from typing import Dict
from typing import Dict, List, Optional, Tuple, Union
from ray.tune.sample import Categorical, Float, Integer, Quantized
from ray.tune.sample import Categorical, Domain, Float, Integer, Quantized
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
@@ -17,8 +17,9 @@ from ray.tune.suggest import Searcher
logger = logging.getLogger(__name__)
def _validate_warmstart(parameter_names, points_to_evaluate,
evaluated_rewards):
def _validate_warmstart(parameter_names: List[str],
points_to_evaluate: List[List],
evaluated_rewards: List):
if points_to_evaluate:
if not isinstance(points_to_evaluate, list):
raise TypeError(
@@ -125,14 +126,14 @@ class SkOptSearch(Searcher):
"""
def __init__(self,
optimizer=None,
space=None,
metric=None,
mode=None,
points_to_evaluate=None,
evaluated_rewards=None,
max_concurrent=None,
use_early_stopped_trials=None):
optimizer: Optional[sko.optimizer.Optimizer] = None,
space: Union[List[str], Dict[str, Union[Tuple, List]]] = None,
metric: Optional[str] = None,
mode: Optional[str] = None,
points_to_evaluate: Optional[List[List]] = None,
evaluated_rewards: Optional[List] = None,
max_concurrent: Optional[int] = None,
use_early_stopped_trials: Optional[bool] = None):
assert sko is not None, """skopt must be installed!
You can install Skopt with the command:
`pip install scikit-optimize`."""
@@ -162,7 +163,7 @@ class SkOptSearch(Searcher):
"names.")
self._parameter_names = space
else:
self._parameter_names = space.keys()
self._parameter_names = list(space.keys())
self._parameter_ranges = space.values()
self._points_to_evaluate = points_to_evaluate
@@ -199,7 +200,8 @@ class SkOptSearch(Searcher):
elif self._mode == "min":
self._metric_op = 1.
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
if self._skopt_opt:
return False
space = self.convert_search_space(config)
@@ -216,7 +218,7 @@ class SkOptSearch(Searcher):
self.setup_skopt()
return True
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
if not self._skopt_opt:
raise RuntimeError(
"Trying to sample a configuration from {}, but no search "
@@ -235,7 +237,10 @@ class SkOptSearch(Searcher):
self._live_trial_mapping[trial_id] = suggested_config
return unflatten_dict(dict(zip(self._parameters, suggested_config)))
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
"""Notification for the completion of trial.
The result is internally negated when interacting with Skopt
@@ -247,24 +252,24 @@ class SkOptSearch(Searcher):
self._process_result(trial_id, result)
self._live_trial_mapping.pop(trial_id)
def _process_result(self, trial_id, result):
def _process_result(self, trial_id: str, result: Dict):
skopt_trial_info = self._live_trial_mapping[trial_id]
self._skopt_opt.tell(skopt_trial_info,
self._metric_op * result[self._metric])
def save(self, checkpoint_path):
def save(self, checkpoint_path: str):
trials_object = (self._initial_points, self._skopt_opt)
with open(checkpoint_path, "wb") as outputFile:
pickle.dump(trials_object, outputFile)
def restore(self, checkpoint_path):
def restore(self, checkpoint_path: str):
with open(checkpoint_path, "rb") as inputFile:
trials_object = pickle.load(inputFile)
self._initial_points = trials_object[0]
self._skopt_opt = trials_object[1]
@staticmethod
def convert_search_space(spec: Dict):
def convert_search_space(spec: Dict) -> Dict:
spec = flatten_dict(spec, prevent_delimiter=True)
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
@@ -273,7 +278,7 @@ class SkOptSearch(Searcher):
"Grid search parameters cannot be automatically converted "
"to a SkOpt search space.")
def resolve_value(domain):
def resolve_value(domain: Domain) -> Union[Tuple, List]:
sampler = domain.get_sampler()
if isinstance(sampler, Quantized):
logger.warning("SkOpt search does not support quantization. "
+36 -24
View File
@@ -2,6 +2,7 @@ import copy
import glob
import logging
import os
from typing import Dict, Optional
from ray.util.debug import log_once
@@ -56,10 +57,10 @@ class Searcher:
CKPT_FILE_TMPL = "searcher-state-{}.pkl"
def __init__(self,
metric=None,
mode=None,
max_concurrent=None,
use_early_stopped_trials=None):
metric: Optional[str] = None,
mode: Optional[str] = None,
max_concurrent: Optional[int] = None,
use_early_stopped_trials: Optional[bool] = None):
if use_early_stopped_trials is False:
raise DeprecationWarning(
"Early stopped trials are now always used. If this is a "
@@ -90,7 +91,8 @@ class Searcher:
else:
raise ValueError("Mode most either be a list or string")
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
"""Pass search properties to searcher.
This method acts as an alternative to instantiating search algorithms
@@ -106,7 +108,7 @@ class Searcher:
"""
return False
def on_trial_result(self, trial_id, result):
def on_trial_result(self, trial_id: str, result: Dict):
"""Optional notification for result during training.
Note that by default, the result dict may include NaNs or
@@ -124,7 +126,10 @@ class Searcher:
"""
pass
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
"""Notification for the completion of trial.
Typically, this method is used for notifying the underlying
@@ -143,7 +148,7 @@ class Searcher:
"""
raise NotImplementedError
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
"""Queries the algorithm to retrieve the next set of parameters.
Arguments:
@@ -159,7 +164,7 @@ class Searcher:
"""
raise NotImplementedError
def save(self, checkpoint_path):
def save(self, checkpoint_path: str):
"""Save state to path for this search algorithm.
Args:
@@ -190,7 +195,7 @@ class Searcher:
"""
raise NotImplementedError
def restore(self, checkpoint_path):
def restore(self, checkpoint_path: str):
"""Restore state for this search algorithm
@@ -213,13 +218,13 @@ class Searcher:
"""
raise NotImplementedError
def get_state(self):
def get_state(self) -> Dict:
raise NotImplementedError
def set_state(self, state):
def set_state(self, state: Dict):
raise NotImplementedError
def save_to_dir(self, checkpoint_dir, session_str="default"):
def save_to_dir(self, checkpoint_dir: str, session_str: str = "default"):
"""Automatically saves the given searcher to the checkpoint_dir.
This is automatically used by tune.run during a Tune job.
@@ -246,7 +251,7 @@ class Searcher:
os.path.join(checkpoint_dir,
self.CKPT_FILE_TMPL.format(session_str)))
def restore_from_dir(self, checkpoint_dir):
def restore_from_dir(self, checkpoint_dir: str):
"""Restores the state of a searcher from a given checkpoint_dir.
Typically, you should use this function to restore from an
@@ -277,12 +282,12 @@ class Searcher:
self.restore(most_recent_checkpoint)
@property
def metric(self):
def metric(self) -> str:
"""The training result objective value attribute."""
return self._metric
@property
def mode(self):
def mode(self) -> str:
"""Specifies if minimizing or maximizing the metric."""
return self._mode
@@ -308,7 +313,10 @@ class ConcurrencyLimiter(Searcher):
tune.run(trainable, search_alg=search_alg)
"""
def __init__(self, searcher, max_concurrent, batch=False):
def __init__(self,
searcher: Searcher,
max_concurrent: int,
batch: bool = False):
assert type(max_concurrent) is int and max_concurrent > 0
self.searcher = searcher
self.max_concurrent = max_concurrent
@@ -318,7 +326,7 @@ class ConcurrencyLimiter(Searcher):
super(ConcurrencyLimiter, self).__init__(
metric=self.searcher.metric, mode=self.searcher.mode)
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
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:
@@ -333,7 +341,10 @@ class ConcurrencyLimiter(Searcher):
self.live_trials.add(trial_id)
return suggestion
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
if trial_id not in self.live_trials:
return
elif self.batch:
@@ -353,19 +364,20 @@ class ConcurrencyLimiter(Searcher):
trial_id, result=result, error=error)
self.live_trials.remove(trial_id)
def get_state(self):
def get_state(self) -> Dict:
state = self.__dict__.copy()
del state["searcher"]
return copy.deepcopy(state)
def set_state(self, state):
def set_state(self, state: Dict):
self.__dict__.update(state)
def on_pause(self, trial_id):
def on_pause(self, trial_id: str):
self.searcher.on_pause(trial_id)
def on_unpause(self, trial_id):
def on_unpause(self, trial_id: str):
self.searcher.on_unpause(trial_id)
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
return self.searcher.set_search_properties(metric, mode, config)
+25 -18
View File
@@ -1,5 +1,7 @@
import copy
import logging
from typing import Any, Dict, Generator, List, Tuple
import numpy
import random
@@ -9,7 +11,8 @@ from ray.tune.sample import Categorical, Domain, Function
logger = logging.getLogger(__name__)
def generate_variants(unresolved_spec):
def generate_variants(
unresolved_spec: Dict) -> Generator[Tuple[Dict, Dict], None, None]:
"""Generates variants from a spec (dict) with unresolved values.
There are two types of unresolved values:
@@ -45,7 +48,7 @@ def generate_variants(unresolved_spec):
yield resolved_vars, spec
def grid_search(values):
def grid_search(values: List) -> Dict[str, List]:
"""Convenience method for specifying grid search over a value.
Arguments:
@@ -63,7 +66,7 @@ _STANDARD_IMPORTS = {
_MAX_RESOLUTION_PASSES = 20
def resolve_nested_dict(nested_dict):
def resolve_nested_dict(nested_dict: Dict) -> Dict[Tuple, Any]:
"""Flattens a nested dict by joining keys into tuple of paths.
Can then be passed into `format_vars`.
@@ -78,7 +81,7 @@ def resolve_nested_dict(nested_dict):
return res
def format_vars(resolved_vars):
def format_vars(resolved_vars: Dict) -> str:
"""Formats the resolved variable dict into a single string."""
out = []
for path, value in sorted(resolved_vars.items()):
@@ -97,7 +100,7 @@ def format_vars(resolved_vars):
return ",".join(out)
def flatten_resolved_vars(resolved_vars):
def flatten_resolved_vars(resolved_vars: Dict) -> Dict:
"""Formats the resolved variable dict into a mapping of (str -> value)."""
flattened_resolved_vars_dict = {}
for pieces, value in resolved_vars.items():
@@ -108,14 +111,15 @@ def flatten_resolved_vars(resolved_vars):
return flattened_resolved_vars_dict
def _clean_value(value):
def _clean_value(value: Any) -> str:
if isinstance(value, float):
return "{:.5}".format(value)
else:
return str(value).replace("/", "_")
def parse_spec_vars(spec):
def parse_spec_vars(spec: Dict) -> Tuple[List[Tuple[Tuple, Any]], List[Tuple[
Tuple, Any]], List[Tuple[Tuple, Any]]]:
resolved, unresolved = _split_resolved_unresolved_values(spec)
resolved_vars = list(resolved.items())
@@ -134,7 +138,7 @@ def parse_spec_vars(spec):
return resolved_vars, domain_vars, grid_vars
def _generate_variants(spec):
def _generate_variants(spec: Dict) -> Tuple[Dict, Dict]:
spec = copy.deepcopy(spec)
_, domain_vars, grid_vars = parse_spec_vars(spec)
@@ -159,19 +163,20 @@ def _generate_variants(spec):
yield resolved_vars, spec
def assign_value(spec, path, value):
def assign_value(spec: Dict, path: Tuple, value: Any):
for k in path[:-1]:
spec = spec[k]
spec[path[-1]] = value
def _get_value(spec, path):
def _get_value(spec: Dict, path: Tuple) -> Any:
for k in path:
spec = spec[k]
return spec
def _resolve_domain_vars(spec, domain_vars):
def _resolve_domain_vars(spec: Dict,
domain_vars: List[Tuple[Tuple, Domain]]) -> Dict:
resolved = {}
error = True
num_passes = 0
@@ -197,7 +202,8 @@ def _resolve_domain_vars(spec, domain_vars):
return resolved
def _grid_search_generator(unresolved_spec, grid_vars):
def _grid_search_generator(unresolved_spec: Dict,
grid_vars: List) -> Generator[Dict, None, None]:
value_indices = [0] * len(grid_vars)
def increment(i):
@@ -225,12 +231,12 @@ def _grid_search_generator(unresolved_spec, grid_vars):
break
def _is_resolved(v):
def _is_resolved(v) -> bool:
resolved, _ = _try_resolve(v)
return resolved
def _try_resolve(v):
def _try_resolve(v) -> Tuple[bool, Any]:
if isinstance(v, Domain):
# Domain to sample from
return False, v
@@ -249,7 +255,8 @@ def _try_resolve(v):
return True, v
def _split_resolved_unresolved_values(spec):
def _split_resolved_unresolved_values(
spec: Dict) -> Tuple[Dict[Tuple, Any], Dict[Tuple, Any]]:
resolved_vars = {}
unresolved_vars = {}
for k, v in spec.items():
@@ -278,11 +285,11 @@ def _split_resolved_unresolved_values(spec):
return resolved_vars, unresolved_vars
def _unresolved_values(spec):
def _unresolved_values(spec: Dict) -> Dict[Tuple, Any]:
return _split_resolved_unresolved_values(spec)[1]
def has_unresolved_values(spec):
def has_unresolved_values(spec: Dict) -> bool:
return True if _unresolved_values(spec) else False
@@ -303,5 +310,5 @@ class _UnresolvedAccessGuard(dict):
class RecursiveDependencyError(Exception):
def __init__(self, msg):
def __init__(self, msg: str):
Exception.__init__(self, msg)
+19 -14
View File
@@ -1,9 +1,10 @@
import copy
import logging
from typing import Dict
from typing import Dict, Optional, Tuple
import ray.cloudpickle as pickle
from ray.tune.sample import Categorical, Float, Integer, Quantized, Uniform
from ray.tune.sample import Categorical, Domain, Float, Integer, Quantized, \
Uniform
from ray.tune.suggest.variant_generator import parse_spec_vars
from ray.tune.utils.util import unflatten_dict
from zoopt import ValueType
@@ -106,11 +107,11 @@ class ZOOptSearch(Searcher):
optimizer = None
def __init__(self,
algo="asracos",
budget=None,
dim_dict=None,
metric=None,
mode=None,
algo: str = "asracos",
budget: Optional[int] = None,
dim_dict: Optional[Dict] = None,
metric: Optional[str] = None,
mode: Optional[str] = None,
**kwargs):
assert zoopt is not None, "Zoopt not found - please install zoopt."
assert budget is not None, "`budget` should not be None!"
@@ -154,7 +155,8 @@ class ZOOptSearch(Searcher):
from zoopt.algos.opt_algorithms.racos.sracos import SRacosTune
self.optimizer = SRacosTune(dimension=dim, parameter=par)
def set_search_properties(self, metric, mode, config):
def set_search_properties(self, metric: Optional[str], mode: Optional[str],
config: Dict) -> bool:
if self._dim_dict:
return False
space = self.convert_search_space(config)
@@ -173,7 +175,7 @@ class ZOOptSearch(Searcher):
self.setup_zoopt()
return True
def suggest(self, trial_id):
def suggest(self, trial_id: str) -> Optional[Dict]:
if not self._dim_dict or not self.optimizer:
raise RuntimeError(
"Trying to sample a configuration from {}, but no search "
@@ -189,7 +191,10 @@ class ZOOptSearch(Searcher):
self._live_trial_mapping[trial_id] = new_trial
return unflatten_dict(new_trial)
def on_trial_complete(self, trial_id, result=None, error=False):
def on_trial_complete(self,
trial_id: str,
result: Optional[Dict] = None,
error: bool = False):
"""Notification for the completion of trial."""
if result:
_solution = self.solution_dict[str(trial_id)]
@@ -200,18 +205,18 @@ class ZOOptSearch(Searcher):
del self._live_trial_mapping[trial_id]
def save(self, checkpoint_path):
def save(self, checkpoint_path: str):
trials_object = self.optimizer
with open(checkpoint_path, "wb") as output:
pickle.dump(trials_object, output)
def restore(self, checkpoint_path):
def restore(self, checkpoint_path: str):
with open(checkpoint_path, "rb") as input:
trials_object = pickle.load(input)
self.optimizer = trials_object
@staticmethod
def convert_search_space(spec: Dict):
def convert_search_space(spec: Dict) -> Dict[str, Tuple]:
spec = copy.deepcopy(spec)
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
@@ -223,7 +228,7 @@ class ZOOptSearch(Searcher):
"Grid search parameters cannot be automatically converted "
"to a ZOOpt search space.")
def resolve_value(domain):
def resolve_value(domain: Domain) -> Tuple:
quantize = None
sampler = domain.get_sampler()