Coverage for mlos_core/mlos_core/optimizers/flaml_optimizer.py: 97%

59 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-20 00:44 +0000

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Contains the :py:class:`.FlamlOptimizer` class. 

7 

8Notes 

9----- 

10See the `Flaml Documentation <https://microsoft.github.io/FLAML/>`_ for more 

11details. 

12""" 

13 

14from typing import Dict, List, NamedTuple, Optional, Union 

15from warnings import warn 

16 

17import ConfigSpace 

18import numpy as np 

19import pandas as pd 

20 

21from mlos_core.data_classes import Observation, Observations, Suggestion 

22from mlos_core.optimizers.optimizer import BaseOptimizer 

23from mlos_core.spaces.adapters.adapter import BaseSpaceAdapter 

24from mlos_core.util import normalize_config 

25 

26 

27class EvaluatedSample(NamedTuple): 

28 """A named tuple representing a sample that has been evaluated.""" 

29 

30 config: dict 

31 score: float 

32 

33 

34class FlamlOptimizer(BaseOptimizer): 

35 """Wrapper class for FLAML Optimizer: A fast library for AutoML and tuning.""" 

36 

37 # The name of an internal objective attribute that is calculated as a weighted 

38 # average of the user provided objective metrics. 

39 _METRIC_NAME = "FLAML_score" 

40 

41 def __init__( 

42 self, 

43 *, # pylint: disable=too-many-arguments 

44 parameter_space: ConfigSpace.ConfigurationSpace, 

45 optimization_targets: List[str], 

46 objective_weights: Optional[List[float]] = None, 

47 space_adapter: Optional[BaseSpaceAdapter] = None, 

48 low_cost_partial_config: Optional[dict] = None, 

49 seed: Optional[int] = None, 

50 ): 

51 """ 

52 Create an MLOS wrapper for FLAML. 

53 

54 Parameters 

55 ---------- 

56 parameter_space : ConfigSpace.ConfigurationSpace 

57 The parameter space to optimize. 

58 

59 optimization_targets : List[str] 

60 The names of the optimization targets to minimize. 

61 

62 objective_weights : Optional[List[float]] 

63 Optional list of weights of optimization targets. 

64 

65 space_adapter : BaseSpaceAdapter 

66 The space adapter class to employ for parameter space transformations. 

67 

68 low_cost_partial_config : dict 

69 A dictionary from a subset of controlled dimensions to the initial low-cost values. 

70 More info: 

71 https://microsoft.github.io/FLAML/docs/FAQ#about-low_cost_partial_config-in-tune 

72 

73 seed : Optional[int] 

74 If provided, calls np.random.seed() with the provided value to set the 

75 seed globally at init. 

76 """ 

77 super().__init__( 

78 parameter_space=parameter_space, 

79 optimization_targets=optimization_targets, 

80 objective_weights=objective_weights, 

81 space_adapter=space_adapter, 

82 ) 

83 

84 # Per upstream documentation, it is recommended to set the seed for 

85 # flaml at the start of its operation globally. 

86 if seed is not None: 

87 np.random.seed(seed) 

88 

89 # pylint: disable=import-outside-toplevel 

90 from mlos_core.spaces.converters.flaml import ( 

91 FlamlDomain, 

92 configspace_to_flaml_space, 

93 ) 

94 

95 self.flaml_parameter_space: Dict[str, FlamlDomain] = configspace_to_flaml_space( 

96 self.optimizer_parameter_space 

97 ) 

98 self.low_cost_partial_config = low_cost_partial_config 

99 

100 self.evaluated_samples: Dict[ConfigSpace.Configuration, EvaluatedSample] = {} 

101 self._suggested_config: Optional[dict] 

102 

103 def _register( 

104 self, 

105 observations: Observations, 

106 ) -> None: 

107 """ 

108 Registers one or more configs/score pairs (observations) with the underlying 

109 optimizer. 

110 

111 Parameters 

112 ---------- 

113 observations : Observations 

114 The set of config/scores to register. 

115 """ 

116 # TODO: Implement bulk registration. 

117 # (e.g., by rebuilding the base optimizer instance with all observations). 

118 for observation in observations: 

119 self._register_single(observation) 

120 

121 def _register_single( 

122 self, 

123 observation: Observation, 

124 ) -> None: 

125 """ 

126 Registers the given config and its score. 

127 

128 Parameters 

129 ---------- 

130 observation : Observation 

131 The observation to register. 

132 """ 

133 if observation.context is not None: 

134 warn( 

135 f"Not Implemented: Ignoring context {list(observation.context.index)}", 

136 UserWarning, 

137 ) 

138 if observation.metadata is not None: 

139 warn( 

140 f"Not Implemented: Ignoring metadata {list(observation.metadata.index)}", 

141 UserWarning, 

142 ) 

143 

144 cs_config: ConfigSpace.Configuration = observation.to_suggestion().to_configspace_config( 

145 self.optimizer_parameter_space 

146 ) 

147 if cs_config in self.evaluated_samples: 

148 warn(f"Configuration {cs_config} was already registered", UserWarning) 

149 self.evaluated_samples[cs_config] = EvaluatedSample( 

150 config=dict(cs_config), 

151 score=float( 

152 np.average(observation.score.astype(float), weights=self._objective_weights) 

153 ), 

154 ) 

155 

156 def _suggest( 

157 self, 

158 *, 

159 context: Optional[pd.Series] = None, 

160 ) -> Suggestion: 

161 """ 

162 Suggests a new configuration. 

163 

164 Sampled at random using ConfigSpace. 

165 

166 Parameters 

167 ---------- 

168 context : None 

169 Not Yet Implemented. 

170 

171 Returns 

172 ------- 

173 suggestion : Suggestion 

174 The suggestion to be evaluated. 

175 """ 

176 if context is not None: 

177 warn(f"Not Implemented: Ignoring context {list(context.index)}", UserWarning) 

178 config: dict = self._get_next_config() 

179 return Suggestion(config=pd.Series(config, dtype=object), context=context, metadata=None) 

180 

181 def register_pending(self, pending: Suggestion) -> None: 

182 raise NotImplementedError() 

183 

184 def _target_function(self, config: dict) -> Union[dict, None]: 

185 """ 

186 Configuration evaluation function called by FLAML optimizer. 

187 

188 FLAML may suggest the same configuration multiple times (due to its 

189 warm-start mechanism). Once FLAML suggests an unseen configuration, we 

190 store it, and stop the optimization process. 

191 

192 Parameters 

193 ---------- 

194 config: dict 

195 Next configuration to be evaluated, as suggested by FLAML. 

196 This config is stored internally and is returned to user, via 

197 `.suggest()` method. 

198 

199 Returns 

200 ------- 

201 result: Union[dict, None] 

202 Dictionary with a single key, `FLAML_score`, if config already 

203 evaluated; `None` otherwise. 

204 """ 

205 cs_config = normalize_config(self.optimizer_parameter_space, config) 

206 if cs_config in self.evaluated_samples: 

207 return {self._METRIC_NAME: self.evaluated_samples[cs_config].score} 

208 

209 self._suggested_config = dict(cs_config) # Cleaned-up version of the config 

210 return None # Returning None stops the process 

211 

212 def _get_next_config(self) -> dict: 

213 """ 

214 Warm-starts a new instance of FLAML, and returns a recommended, unseen new 

215 configuration. 

216 

217 Since FLAML does not provide an ask-and-tell interface, we need to create a 

218 new instance of FLAML each time we get asked for a new suggestion. This is 

219 suboptimal performance-wise, but works. 

220 To do so, we use any previously evaluated configs to bootstrap FLAML (i.e., 

221 warm-start). 

222 For more info: 

223 https://microsoft.github.io/FLAML/docs/Use-Cases/Tune-User-Defined-Function#warm-start 

224 

225 Returns 

226 ------- 

227 result: dict 

228 A dictionary with a single key that is equal to the name of the optimization target, 

229 if config already evaluated; `None` otherwise. 

230 

231 Raises 

232 ------ 

233 RuntimeError: if FLAML did not suggest a previously unseen configuration. 

234 """ 

235 from flaml import tune # pylint: disable=import-outside-toplevel 

236 

237 # Parse evaluated configs to format used by FLAML 

238 points_to_evaluate: list = [] 

239 evaluated_rewards: list = [] 

240 if len(self.evaluated_samples) > 0: 

241 points_to_evaluate = [ 

242 dict(normalize_config(self.optimizer_parameter_space, conf)) 

243 for conf in self.evaluated_samples 

244 ] 

245 evaluated_rewards = [s.score for s in self.evaluated_samples.values()] 

246 

247 # Warm start FLAML optimizer 

248 self._suggested_config = None 

249 tune.run( 

250 self._target_function, 

251 config=self.flaml_parameter_space, 

252 mode="min", 

253 metric=self._METRIC_NAME, 

254 points_to_evaluate=points_to_evaluate, 

255 evaluated_rewards=evaluated_rewards, 

256 num_samples=len(points_to_evaluate) + 1, 

257 low_cost_partial_config=self.low_cost_partial_config, 

258 verbose=0, 

259 ) 

260 if self._suggested_config is None: 

261 raise RuntimeError("FLAML did not produce a suggestion") 

262 

263 return self._suggested_config # type: ignore[unreachable]