dataclasses for mlos_core Optimizer return values

mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py

@@ -22,6 +22,7 @@
 from mlos_bench.services.base_service import Service
 from mlos_bench.tunables.tunable import TunableValue
 from mlos_bench.tunables.tunable_groups import TunableGroups
+from mlos_core.optimizers.observations import Observation
 from mlos_core.optimizers import (
     DEFAULT_OPTIMIZER_TYPE,
     BaseOptimizer,
@@ -128,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(configs=df_configs, scores=df_scores)
+        self._opt.register(observation=Observation(config=df_configs, performance=df_scores))
 
         if _LOG.isEnabledFor(logging.DEBUG):
             (score, _) = self.get_best_observation()
@@ -198,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,
@@ -221,15 +224,19 @@ def register(
             # TODO: Specify (in the config) which metrics to pass to the optimizer.
             # Issue: https://github.com/microsoft/MLOS/issues/745
             self._opt.register(
-                configs=df_config,
-                scores=pd.DataFrame([registered_score], dtype=float),
+                observation=Observation(
+                    config=df_config,
+                    performance=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()
+        (df_config, df_score, _df_context, _metadata) = (
+            self._opt.get_best_observations().to_legacy()
+        )
         if len(df_config) == 0:
             return (None, None)
         params = configspace_data_to_tunable_values(df_config.iloc[0].to_dict())

mlos_core/mlos_core/optimizers/bayesian_optimizers/smac_optimizer.py

@@ -18,6 +18,7 @@
 import numpy.typing as npt
 import pandas as pd
 
+from mlos_core.optimizers.observations import Observation, Suggestion
 from mlos_core.optimizers.bayesian_optimizers.bayesian_optimizer import (
     BaseBayesianOptimizer,
 )
@@ -272,42 +273,32 @@ def _dummy_target_func(config: ConfigSpace.Configuration, seed: int = 0) -> None
     def _register(
         self,
         *,
-        configs: pd.DataFrame,
-        scores: pd.DataFrame,
-        context: Optional[pd.DataFrame] = None,
-        metadata: Optional[pd.DataFrame] = None,
+        observation: Observation,
     ) -> None:
         """
         Registers the given configs and scores.
 
         Parameters
         ----------
-        configs : pd.DataFrame
-            Dataframe of configs / parameters. The columns are parameter names and
-            the rows are the configs.
-
-        scores : pd.DataFrame
-            Scores from running the configs. The index is the same as the index of
-            the configs.
-
-        context : pd.DataFrame
-            Not Yet Implemented.
-
-        metadata: pd.DataFrame
-            Not Yet Implemented.
+        observation: Observation
+            The observation to register.
         """
         from smac.runhistory import (  # pylint: disable=import-outside-toplevel
             StatusType,
             TrialInfo,
             TrialValue,
         )
 
-        if context is not None:
-            warn(f"Not Implemented: Ignoring context {list(context.columns)}", UserWarning)
+        if observation.context is not None:
+            warn(
+                f"Not Implemented: Ignoring context {list(observation.context.columns)}",
+                UserWarning,
+            )
 
         # Register each trial (one-by-one)
         for config, (_i, score) in zip(
-            self._to_configspace_configs(configs=configs), scores.iterrows()
+            self._to_configspace_configs(configs=observation.config),
+            observation.performance.iterrows(),
         ):
             # Retrieve previously generated TrialInfo (returned by .ask()) or create
             # new TrialInfo instance
@@ -325,7 +316,7 @@ def _suggest(
         self,
         *,
         context: Optional[pd.DataFrame] = None,
-    ) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:
+    ) -> Suggestion:
         """
         Suggests a new configuration.
 
@@ -357,15 +348,9 @@ 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, *, suggestion: Suggestion) -> None:
         raise NotImplementedError()
 
     def surrogate_predict(
@@ -383,11 +368,14 @@ def surrogate_predict(
             raise NotImplementedError("Space adapter not supported for surrogate_predict.")
 
         # pylint: disable=protected-access
-        if len(self._observations) <= self.base_optimizer._initial_design._n_configs:
+        if (
+            sum(len(o.config.index) for o in self._observations)
+            <= self.base_optimizer._initial_design._n_configs
+        ):
             raise RuntimeError(
                 "Surrogate model can make predictions *only* after "
                 "all initial points have been evaluated "
-                f"{len(self._observations)} <= {self.base_optimizer._initial_design._n_configs}"
+                f"{sum(len(o.config.index) for o in self._observations)} <= {self.base_optimizer._initial_design._n_configs}"
             )
         if self.base_optimizer._config_selector._model is None:
             raise RuntimeError("Surrogate model is not yet trained")
@@ -440,7 +428,8 @@ def _to_configspace_configs(self, *, configs: pd.DataFrame) -> List[ConfigSpace.
         configs : list
             List of ConfigSpace configs.
         """
+        values = [config.to_dict() for (_, config) in configs.astype("O").iterrows()]
         return [
-            ConfigSpace.Configuration(self.optimizer_parameter_space, values=config.to_dict())
-            for (_, config) in configs.astype("O").iterrows()
+            ConfigSpace.Configuration(self.optimizer_parameter_space, values=value)
+            for value in values
         ]

mlos_core/mlos_core/optimizers/flaml_optimizer.py

@@ -11,6 +11,7 @@
 import numpy as np
 import pandas as pd
 
+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
@@ -92,38 +93,29 @@ def __init__(
         self.evaluated_samples: Dict[ConfigSpace.Configuration, EvaluatedSample] = {}
         self._suggested_config: Optional[dict]
 
-    def _register(
-        self,
-        *,
-        configs: pd.DataFrame,
-        scores: pd.DataFrame,
-        context: Optional[pd.DataFrame] = None,
-        metadata: Optional[pd.DataFrame] = None,
-    ) -> None:
+    def _register(self, *, observation: Observation) -> None:
         """
         Registers the given configs and scores.
 
         Parameters
         ----------
-        configs : pd.DataFrame
-            Dataframe of configs / parameters. The columns are parameter names and
-            the rows are the configs.
-
-        scores : pd.DataFrame
-            Scores from running the configs. The index is the same as the index of the configs.
-
-        context : None
-            Not Yet Implemented.
-
-        metadata : None
-            Not Yet Implemented.
+        observation: Observation
+            The observation to register.
         """
-        if context is not None:
-            warn(f"Not Implemented: Ignoring context {list(context.columns)}", UserWarning)
-        if metadata is not None:
-            warn(f"Not Implemented: Ignoring metadata {list(metadata.columns)}", UserWarning)
+        if observation.context is not None:
+            warn(
+                f"Not Implemented: Ignoring context {list(observation.context.columns)}",
+                UserWarning,
+            )
+        if observation.metadata is not None:
+            warn(
+                f"Not Implemented: Ignoring metadata {list(observation.metadata.columns)}",
+                UserWarning,
+            )
 
-        for (_, config), (_, score) in zip(configs.astype("O").iterrows(), scores.iterrows()):
+        for (_, config), (_, score) in zip(
+            observation.config.astype("O").iterrows(), observation.performance.iterrows()
+        ):
             cs_config: ConfigSpace.Configuration = ConfigSpace.Configuration(
                 self.optimizer_parameter_space, values=config.to_dict()
             )
@@ -138,7 +130,7 @@ def _suggest(
         self,
         *,
         context: Optional[pd.DataFrame] = None,
-    ) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:
+    ) -> Suggestion:
         """
         Suggests a new configuration.
 
@@ -151,24 +143,15 @@ def _suggest(
 
         Returns
         -------
-        configuration : pd.DataFrame
-            Pandas dataframe with a single row. Column names are the parameter names.
-
-        metadata : None
-            Not implemented.
+        suggestion: Suggestion
+            The suggestion to evaluate.
         """
         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)
 
-    def register_pending(
-        self,
-        *,
-        configs: pd.DataFrame,
-        context: Optional[pd.DataFrame] = None,
-        metadata: Optional[pd.DataFrame] = None,
-    ) -> None:
+    def register_pending(self, *, suggestion: Suggestion) -> None:
         raise NotImplementedError()
 
     def _target_function(self, config: dict) -> Union[dict, None]: