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

59 statements  

« prev     ^ index     » next       coverage.py v7.6.10, created at 2025-01-21 01:50 +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 NamedTuple 

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: list[float] | None = None, 

47 space_adapter: BaseSpaceAdapter | None = None, 

48 low_cost_partial_config: dict | None = None, 

49 seed: int | None = 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 : int | None 

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: dict | None 

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: pd.Series | None = 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) -> 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: 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]