fix test cases for introducing tuple classes

mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py

@@ -129,7 +129,7 @@ def bulk_register(
 
  # TODO: Specify (in the config) which metrics to pass to the optimizer.
  # Issue: https://github.com/microsoft/MLOS/issues/745
- self._opt.register(observation=Observation(config=df_configs, performance=df_scores))
+ self._opt.register(observation=Observation(config=df_configs, score=df_scores))
 
  if _LOG.isEnabledFor(logging.DEBUG):
  (score, _) = self.get_best_observation()
@@ -199,10 +199,12 @@ def suggest(self) -> TunableGroups:
  tunables = super().suggest()
  if self._start_with_defaults:
  _LOG.info("Use default values for the first trial")
- df_config, _metadata = self._opt.suggest(defaults=self._start_with_defaults)
+ suggestion = self._opt.suggest(defaults=self._start_with_defaults)
  self._start_with_defaults = False
- _LOG.info("Iteration %d :: Suggest:\n%s", self._iter, df_config)
- return tunables.assign(configspace_data_to_tunable_values(df_config.loc[0].to_dict()))
+ _LOG.info("Iteration %d :: Suggest:\n%s", self._iter, suggestion.config)
+ return tunables.assign(
+ configspace_data_to_tunable_values(suggestion.config.loc[0].to_dict())
+ )
 
  def register(
  self,
@@ -224,18 +226,18 @@ def register(
  self._opt.register(
  observation=Observation(
  config=df_config,
- performance=pd.DataFrame([registered_score], dtype=float),
+ score=pd.DataFrame([registered_score], dtype=float),
  )
  )
  return registered_score
 
  def get_best_observation(
  self,
  ) -> Union[Tuple[Dict[str, float], TunableGroups], Tuple[None, None]]:
- (df_config, df_score, _df_context) = self._opt.get_best_observations()
- if len(df_config) == 0:
+ best_observations = self._opt.get_best_observations()
+ if len(best_observations.config.index) == 0:
  return (None, None)
- params = configspace_data_to_tunable_values(df_config.iloc[0].to_dict())
- scores = self._adjust_signs_df(df_score).iloc[0].to_dict()
+ params = configspace_data_to_tunable_values(best_observations.config.iloc[0].to_dict())
+ scores = self._adjust_signs_df(best_observations.score).iloc[0].to_dict()
  _LOG.debug("Best observation: %s score: %s", params, scores)
  return (scores, self._tunables.copy().assign(params))

mlos_bench/mlos_bench/tests/optimizers/toy_optimization_loop_test.py

@@ -7,6 +7,7 @@
 import logging
 from typing import Tuple
 
+from mlos_core.optimizers.observations import Suggestion
 import pytest
 
 from mlos_bench.environments.base_environment import Environment
@@ -43,7 +44,10 @@ def _optimize(env: Environment, opt: Optimizer) -> Tuple[float, TunableGroups]:
  config_df = config_to_dataframe(config)
  logger("config: %s", str(config))
  try:
- logger("prediction: %s", opt._opt.surrogate_predict(configs=config_df))
+ logger(
+ "prediction: %s",
+ opt._opt.surrogate_predict(suggestion=Suggestion(config=config_df)),
+ )
  except RuntimeError:
  pass
 

mlos_core/mlos_core/optimizers/bayesian_optimizers/bayesian_optimizer.py

@@ -7,6 +7,7 @@
 from abc import ABCMeta, abstractmethod
 from typing import Optional
 
+from mlos_core.optimizers.observations import Suggestion
 import numpy.typing as npt
 import pandas as pd
 
@@ -17,45 +18,27 @@ class BaseBayesianOptimizer(BaseOptimizer, metaclass=ABCMeta):
  """Abstract base class defining the interface for Bayesian optimization."""
 
  @abstractmethod
- def surrogate_predict(
- self,
- *,
- configs: pd.DataFrame,
- context: Optional[pd.DataFrame] = None,
- ) -> npt.NDArray:
+ def surrogate_predict(self, *, suggestion: Suggestion) -> npt.NDArray:
  """
  Obtain a prediction from this Bayesian optimizer's surrogate model for the given
  configuration(s).
 
  Parameters
  ----------
- configs : pd.DataFrame
- Dataframe of configs / parameters. The columns are parameter names and
- the rows are the configs.
-
- context : pd.DataFrame
- Not Yet Implemented.
+ suggestion: Suggestion
+ The suggestion containing the configuration(s) to predict.
  """
  pass # pylint: disable=unnecessary-pass # pragma: no cover
 
  @abstractmethod
- def acquisition_function(
- self,
- *,
- configs: pd.DataFrame,
- context: Optional[pd.DataFrame] = None,
- ) -> npt.NDArray:
+ def acquisition_function(self, *, suggestion: Suggestion) -> npt.NDArray:
  """
  Invokes the acquisition function from this Bayesian optimizer for the given
  configuration.
 
  Parameters
  ----------
- configs : pd.DataFrame
- Dataframe of configs / parameters. The columns are parameter names and
- the rows are the configs.
-
- context : pd.DataFrame
- Not Yet Implemented.
+ suggestion: Suggestion
+ The suggestion containing the configuration(s) to evaluate.
  """
  pass # pylint: disable=unnecessary-pass # pragma: no cover

mlos_core/mlos_core/optimizers/bayesian_optimizers/smac_optimizer.py

@@ -18,7 +18,7 @@
 import numpy.typing as npt
 import pandas as pd
 
-from mlos_core.optimizers.observations import Observation
+from mlos_core.optimizers.observations import Observation, Suggestion
 from mlos_core.optimizers.bayesian_optimizers.bayesian_optimizer import (
  BaseBayesianOptimizer,
 )
@@ -294,7 +294,7 @@ def _register(self, *, observation: Observation) -> None:
  # Register each trial (one-by-one)
  for config, (_i, score) in zip(
  self._to_configspace_configs(configs=observation.config),
- observation.performance.iterrows(),
+ observation.score.iterrows(),
  ):
  # Retrieve previously generated TrialInfo (returned by .ask()) or create
  # new TrialInfo instance
@@ -312,7 +312,7 @@ def _suggest(
  self,
  *,
  context: Optional[pd.DataFrame] = None,
- ) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:
+ ) -> Suggestion:
  """
  Suggests a new configuration.
 
@@ -323,11 +323,8 @@ def _suggest(
 
  Returns
  -------
- configuration : pd.DataFrame
- Pandas dataframe with a single row. Column names are the parameter names.
-
- metadata : Optional[pd.DataFrame]
- Not yet implemented.
+ suggestion: Suggestion
+ The suggestion to evaluate.
  """
  if TYPE_CHECKING:
  # pylint: disable=import-outside-toplevel,unused-import
@@ -344,28 +341,20 @@ def _suggest(
  config_df = pd.DataFrame(
  [trial.config], columns=list(self.optimizer_parameter_space.keys())
  )
- return config_df, None
+ return Suggestion(config=config_df, context=context, metadata=None)
 
- def register_pending(
- self,
- *,
- configs: pd.DataFrame,
- context: Optional[pd.DataFrame] = None,
- metadata: Optional[pd.DataFrame] = None,
- ) -> None:
+ def register_pending(self, *, pending: Suggestion) -> None:
  raise NotImplementedError()
 
- def surrogate_predict(
- self,
- *,
- configs: pd.DataFrame,
- context: Optional[pd.DataFrame] = None,
- ) -> npt.NDArray:
+ def surrogate_predict(self, *, suggestion: Suggestion) -> npt.NDArray:
  # pylint: disable=import-outside-toplevel
  from smac.utils.configspace import convert_configurations_to_array
 
- if context is not None:
- warn(f"Not Implemented: Ignoring context {list(context.columns)}", UserWarning)
+ if suggestion.context is not None:
+ warn(
+ f"Not Implemented: Ignoring context {list(suggestion.context.columns)}",
+ UserWarning,
+ )
  if self._space_adapter and not isinstance(self._space_adapter, IdentityAdapter):
  raise NotImplementedError("Space adapter not supported for surrogate_predict.")
 
@@ -380,29 +369,27 @@ def surrogate_predict(
  raise RuntimeError("Surrogate model is not yet trained")
 
  config_array: npt.NDArray = convert_configurations_to_array(
- self._to_configspace_configs(configs=configs)
+ self._to_configspace_configs(configs=suggestion.config)
  )
  mean_predictions, _ = self.base_optimizer._config_selector._model.predict(config_array)
  return mean_predictions.reshape(
  -1,
  )
 
- def acquisition_function(
- self,
- *,
- configs: pd.DataFrame,
- context: Optional[pd.DataFrame] = None,
- ) -> npt.NDArray:
- if context is not None:
- warn(f"Not Implemented: Ignoring context {list(context.columns)}", UserWarning)
+ def acquisition_function(self, *, suggestion: Suggestion) -> npt.NDArray:
+ if suggestion.context is not None:
+ warn(
+ f"Not Implemented: Ignoring context {list(suggestion.context.columns)}",
+ UserWarning,
+ )
  if self._space_adapter:
  raise NotImplementedError()
 
  # pylint: disable=protected-access
  if self.base_optimizer._config_selector._acquisition_function is None:
  raise RuntimeError("Acquisition function is not yet initialized")
 
- cs_configs: list = self._to_configspace_configs(configs=configs)
+ cs_configs: list = self._to_configspace_configs(configs=suggestion.config)
  return self.base_optimizer._config_selector._acquisition_function(cs_configs).reshape(
  -1,
  )
@@ -427,8 +414,7 @@ def _to_configspace_configs(self, *, configs: pd.DataFrame) -> List[ConfigSpace.
  configs : list
  List of ConfigSpace configs.
  """
- assert isinstance(configs, pd.DataFrame)
  return [
  ConfigSpace.Configuration(self.optimizer_parameter_space, values=config.to_dict())
- for (_, config) in configs.iterrows()
+ for (_, config) in configs.astype("O").iterrows()
  ]

mlos_core/mlos_core/optimizers/flaml_optimizer.py

@@ -11,7 +11,7 @@
 import numpy as np
 import pandas as pd
 
-from mlos_core.optimizers.observations import Observation
+from mlos_core.optimizers.observations import Observation, Suggestion
 from mlos_core.optimizers.optimizer import BaseOptimizer
 from mlos_core.spaces.adapters.adapter import BaseSpaceAdapter
 from mlos_core.util import normalize_config
@@ -124,9 +124,9 @@ def _register(self, *, observation: Observation) -> None:
  )
 
  assert isinstance(observation.config, pd.DataFrame)
- assert isinstance(observation.performance, pd.DataFrame)
+ assert isinstance(observation.score, pd.DataFrame)
  for (_, config), (_, score) in zip(
- observation.config.iterrows(), observation.performance.iterrows()
+ observation.config.astype("O").iterrows(), observation.score.iterrows()
  ):
  cs_config: ConfigSpace.Configuration = ConfigSpace.Configuration(
  self.optimizer_parameter_space, values=config.to_dict()
@@ -142,7 +142,7 @@ def _suggest(
  self,
  *,
  context: Optional[pd.DataFrame] = None,
- ) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:
+ ) -> Suggestion:
  """
  Suggests a new configuration.
 
@@ -164,15 +164,9 @@ def _suggest(
  if context is not None:
  warn(f"Not Implemented: Ignoring context {list(context.columns)}", UserWarning)
  config: dict = self._get_next_config()
- return pd.DataFrame(config, index=[0]), None
+ return Suggestion(config=pd.DataFrame(config, index=[0]), context=context, metadata=None)
 
- def register_pending(
- self,
- *,
- configs: pd.DataFrame,
- context: Optional[pd.DataFrame] = None,
- metadata: Optional[pd.DataFrame] = None,
- ) -> None:
+ def register_pending(self, *, pending: Suggestion) -> None:
  raise NotImplementedError()
 
  def _target_function(self, config: dict) -> Union[dict, None]: