From fdcc3e3ef610b94be8c75845d72c0a7c0b963cc2 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Wed, 10 Jan 2024 21:44:08 +0000
Subject: [PATCH 01/24] Add support for storing the optimization target and
 direction of an experiment

---
 mlos_bench/mlos_bench/optimizers/base_optimizer.py | 7 +++++++
 mlos_bench/mlos_bench/run.py                       | 3 ++-
 mlos_bench/mlos_bench/storage/base_storage.py      | 9 +++++++--
 mlos_bench/mlos_bench/storage/sql/experiment.py    | 6 +++++-
 mlos_bench/mlos_bench/storage/sql/storage.py       | 4 +++-
 mlos_bench/mlos_bench/storage/sql/trial.py         | 4 +++-
 mlos_bench/mlos_bench/tests/storage/conftest.py    | 1 +
 7 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/mlos_bench/mlos_bench/optimizers/base_optimizer.py b/mlos_bench/mlos_bench/optimizers/base_optimizer.py
index 62b24e642d6..663bef5e7da 100644
--- a/mlos_bench/mlos_bench/optimizers/base_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/base_optimizer.py
@@ -167,6 +167,13 @@ def target(self) -> str:
         """
         return self._opt_target
 
+    @property
+    def direction(self) -> str:
+        """
+        The direction to optimize the target metric (e.g., min or max).
+        """
+        return 'min' if self.is_min else 'max'
+
     @property
     def supports_preload(self) -> bool:
         """
diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index 70a416e4b04..a9c2652ec6d 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -86,7 +86,8 @@ def _optimize(*,
             trial_id=trial_id,
             root_env_config=root_env_config,
             description=env.name,
-            opt_target=opt.target
+            opt_target=opt.target,
+            opt_direction=opt.direction,
          ) as exp:
 
         _LOG.info("Experiment: %s Env: %s Optimizer: %s", exp, env, opt)
diff --git a/mlos_bench/mlos_bench/storage/base_storage.py b/mlos_bench/mlos_bench/storage/base_storage.py
index e5000b5d6fd..39ec9e5a65c 100644
--- a/mlos_bench/mlos_bench/storage/base_storage.py
+++ b/mlos_bench/mlos_bench/storage/base_storage.py
@@ -83,7 +83,8 @@ def experiment(self, *,
                    trial_id: int,
                    root_env_config: str,
                    description: str,
-                   opt_target: str) -> 'Storage.Experiment':
+                   opt_target: str,
+                   opt_direction: str) -> 'Storage.Experiment':
         """
         Create a new experiment in the storage.
 
@@ -103,6 +104,8 @@ def experiment(self, *,
             Human-readable description of the experiment.
         opt_target : str
             Name of metric we're optimizing for.
+        opt_direction: str
+            Direction to optimize the metric (e.g., min or max)
 
         Returns
         -------
@@ -249,12 +252,14 @@ class Trial(metaclass=ABCMeta):
 
         def __init__(self, *,
                      tunables: TunableGroups, experiment_id: str, trial_id: int,
-                     config_id: int, opt_target: str, config: Optional[Dict[str, Any]] = None):
+                     config_id: int, opt_target: str, opt_direction: str,
+                     config: Optional[Dict[str, Any]] = None):
             self._tunables = tunables
             self._experiment_id = experiment_id
             self._trial_id = trial_id
             self._config_id = config_id
             self._opt_target = opt_target
+            self._opt_direction = opt_direction
             self._config = config or {}
 
         def __repr__(self) -> str:
diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index 8a271a95562..4127a8d82a7 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -35,13 +35,15 @@ def __init__(self, *,
                  trial_id: int,
                  root_env_config: str,
                  description: str,
-                 opt_target: str):
+                 opt_target: str,
+                 opt_direction: str):
         super().__init__(tunables, experiment_id, root_env_config)
         self._engine = engine
         self._schema = schema
         self._trial_id = trial_id
         self._description = description
         self._opt_target = opt_target
+        self._opt_direction = opt_direction
 
     def _setup(self) -> None:
         super()._setup()
@@ -180,6 +182,7 @@ def pending_trials(self) -> Iterator[Storage.Trial]:
                     trial_id=trial.trial_id,
                     config_id=trial.config_id,
                     opt_target=self._opt_target,
+                    opt_direction=self._opt_direction,
                     config=config,
                 )
 
@@ -228,6 +231,7 @@ def new_trial(self, tunables: TunableGroups,
                     trial_id=self._trial_id,
                     config_id=config_id,
                     opt_target=self._opt_target,
+                    opt_direction=self._opt_direction,
                     config=config,
                 )
                 self._trial_id += 1
diff --git a/mlos_bench/mlos_bench/storage/sql/storage.py b/mlos_bench/mlos_bench/storage/sql/storage.py
index 32d04e7d442..64543173c36 100644
--- a/mlos_bench/mlos_bench/storage/sql/storage.py
+++ b/mlos_bench/mlos_bench/storage/sql/storage.py
@@ -62,7 +62,8 @@ def experiment(self, *,
                    trial_id: int,
                    root_env_config: str,
                    description: str,
-                   opt_target: str) -> Storage.Experiment:
+                   opt_target: str,
+                   opt_direction: str) -> Storage.Experiment:
         return Experiment(
             engine=self._engine,
             schema=self._schema,
@@ -72,6 +73,7 @@ def experiment(self, *,
             root_env_config=root_env_config,
             description=description,
             opt_target=opt_target,
+            opt_direction=opt_direction,
         )
 
     @property
diff --git a/mlos_bench/mlos_bench/storage/sql/trial.py b/mlos_bench/mlos_bench/storage/sql/trial.py
index 242ab5f0e94..2b6ada313e5 100644
--- a/mlos_bench/mlos_bench/storage/sql/trial.py
+++ b/mlos_bench/mlos_bench/storage/sql/trial.py
@@ -29,13 +29,15 @@ class Trial(Storage.Trial):
     def __init__(self, *,
                  engine: Engine, schema: DbSchema, tunables: TunableGroups,
                  experiment_id: str, trial_id: int, config_id: int,
-                 opt_target: str, config: Optional[Dict[str, Any]] = None):
+                 opt_target: str, opt_direction: str,
+                 config: Optional[Dict[str, Any]] = None):
         super().__init__(
             tunables=tunables,
             experiment_id=experiment_id,
             trial_id=trial_id,
             config_id=config_id,
             opt_target=opt_target,
+            opt_direction=opt_direction,
             config=config,
         )
         self._engine = engine
diff --git a/mlos_bench/mlos_bench/tests/storage/conftest.py b/mlos_bench/mlos_bench/tests/storage/conftest.py
index d3de677142a..eeb5626f0de 100644
--- a/mlos_bench/mlos_bench/tests/storage/conftest.py
+++ b/mlos_bench/mlos_bench/tests/storage/conftest.py
@@ -33,4 +33,5 @@ def exp_storage_memory_sql(tunable_groups: TunableGroups) -> Storage.Experiment:
         root_env_config="environment.jsonc",
         description="pytest experiment",
         opt_target="score",
+        opt_direction="min",
     ).__enter__()

From cf696740d357f3b7427b55e6b5c48c301117f38a Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 18:40:26 +0000
Subject: [PATCH 02/24] Rename and improve tunable vs config API and
 documentation

---
 mlos_bench/mlos_bench/run.py                       |  2 +-
 .../mlos_bench/storage/base_experiment_data.py     |  2 +-
 mlos_bench/mlos_bench/storage/base_storage.py      | 11 +++++++++--
 mlos_bench/mlos_bench/storage/base_trial_data.py   | 14 +++++++++-----
 mlos_bench/mlos_bench/storage/sql/experiment.py    |  6 +++++-
 mlos_bench/mlos_bench/storage/sql/trial_data.py    | 10 +++++++---
 6 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index a9c2652ec6d..a0378e97e7a 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -110,7 +110,7 @@ def _optimize(*,
             tunables = opt_context.suggest()
 
             if config_id > 0:
-                tunable_values = exp.load_config(config_id)
+                tunable_values = exp.load_tunable_config(config_id)
                 tunables.assign(tunable_values)
                 _LOG.info("Load config from storage: %d", config_id)
                 if _LOG.isEnabledFor(logging.DEBUG):
diff --git a/mlos_bench/mlos_bench/storage/base_experiment_data.py b/mlos_bench/mlos_bench/storage/base_experiment_data.py
index 22025f290ae..f9240d8fe7a 100644
--- a/mlos_bench/mlos_bench/storage/base_experiment_data.py
+++ b/mlos_bench/mlos_bench/storage/base_experiment_data.py
@@ -79,6 +79,6 @@ def results(self) -> pandas.DataFrame:
         results : pandas.DataFrame
             A DataFrame with configurations and results from all trials of the experiment.
             Has columns [trial_id, config_id, ts_start, ts_end, status]
-            followed by config parameters and trial results. The latter can be NULLs
+            followed by tunable config parameters and trial results. The latter can be NULLs
             if the trial was not successful.
         """
diff --git a/mlos_bench/mlos_bench/storage/base_storage.py b/mlos_bench/mlos_bench/storage/base_storage.py
index 39ec9e5a65c..282d64a429f 100644
--- a/mlos_bench/mlos_bench/storage/base_storage.py
+++ b/mlos_bench/mlos_bench/storage/base_storage.py
@@ -188,7 +188,7 @@ def merge(self, experiment_ids: List[str]) -> None:
             """
 
         @abstractmethod
-        def load_config(self, config_id: int) -> Dict[str, Any]:
+        def load_tunable_config(self, config_id: int) -> Dict[str, Any]:
             """
             Load tunable values for a given config ID.
             """
@@ -282,7 +282,9 @@ def config_id(self) -> int:
         @property
         def tunables(self) -> TunableGroups:
             """
-            Tunable parameters of the current trial.
+            Tunable parameters of the current trial
+
+            (e.g., application Environment's "config")
             """
             return self._tunables
 
@@ -290,6 +292,11 @@ def config(self, global_config: Optional[Dict[str, Any]] = None) -> Dict[str, An
             """
             Produce a copy of the global configuration updated
             with the parameters of the current trial.
+
+            Note: this is not the target Environment's "config" (i.e., tunable
+            params), but rather the internal "config" which consists of a
+            combination of somewhat more static variables defined in the json config
+            files.
             """
             config = self._config.copy()
             config.update(global_config or {})
diff --git a/mlos_bench/mlos_bench/storage/base_trial_data.py b/mlos_bench/mlos_bench/storage/base_trial_data.py
index 58d0c9bb43e..71ad7e75a42 100644
--- a/mlos_bench/mlos_bench/storage/base_trial_data.py
+++ b/mlos_bench/mlos_bench/storage/base_trial_data.py
@@ -80,14 +80,16 @@ def status(self) -> Status:
 
     @property
     @abstractmethod
-    def config(self) -> pandas.DataFrame:
+    def tunable_config(self) -> pandas.DataFrame:
         """
-        Retrieve the trials' configuration from the storage.
+        Retrieve the trials' tunable configuration from the storage.
+
+        Note: this corresponds to the Trial object's "tunables" property.
 
         Returns
         -------
         config : pandas.DataFrame
-            A dataframe with the configuration of the current trial.
+            A dataframe with the tunable configuration of the current trial.
             It has two `str` columns, "parameter" and "value".
         """
 
@@ -124,11 +126,13 @@ def telemetry(self) -> pandas.DataFrame:
     @abstractmethod
     def metadata(self) -> pandas.DataFrame:
         """
-        Retrieve the trials' metadata.
+        Retrieve the trials' metadata parameters.
+
+        Note: this corresponds to the Trial object's "config" property.
 
         Returns
         -------
-        config : pandas.DataFrame
+        metadata : pandas.DataFrame
             An optional dataframe with the metadata associated with the trial.
             It has two `str` columns, "parameter" and "value".
             Returns an empty dataframe if there is no metadata.
diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index 4127a8d82a7..2338161b3b8 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -91,7 +91,7 @@ def merge(self, experiment_ids: List[str]) -> None:
         _LOG.info("Merge: %s <- %s", self._experiment_id, experiment_ids)
         raise NotImplementedError()
 
-    def load_config(self, config_id: int) -> Dict[str, Any]:
+    def load_tunable_config(self, config_id: int) -> Dict[str, Any]:
         with self._engine.connect() as conn:
             return self._get_params(conn, self._schema.config_param, config_id=config_id)
 
@@ -219,10 +219,14 @@ def new_trial(self, tunables: TunableGroups,
                     ts_start=datetime.utcnow(),
                     status='PENDING',
                 ))
+
+                # Note: config here is the framework config, not the target
+                # environment config (i.e., tunables).
                 if config is not None:
                     self._save_params(
                         conn, self._schema.trial_param, config,
                         exp_id=self._experiment_id, trial_id=self._trial_id)
+
                 trial = Trial(
                     engine=self._engine,
                     schema=self._schema,
diff --git a/mlos_bench/mlos_bench/storage/sql/trial_data.py b/mlos_bench/mlos_bench/storage/sql/trial_data.py
index 75d6d05e21e..8a0dd187b47 100644
--- a/mlos_bench/mlos_bench/storage/sql/trial_data.py
+++ b/mlos_bench/mlos_bench/storage/sql/trial_data.py
@@ -42,9 +42,11 @@ def __init__(self, *,
         self._schema = schema
 
     @property
-    def config(self) -> pandas.DataFrame:
+    def tunable_config(self) -> pandas.DataFrame:
         """
-        Retrieve the trials' configuration from the storage.
+        Retrieve the trials' tunable configuration from the storage.
+
+        Note: this corresponds to the Trial object's "tunables" property.
         """
         with self._engine.connect() as conn:
             cur_config = conn.execute(
@@ -98,7 +100,9 @@ def telemetry(self) -> pandas.DataFrame:
     @property
     def metadata(self) -> pandas.DataFrame:
         """
-        Retrieve the trials' metadata.
+        Retrieve the trials' metadata params.
+
+        Note: this corresponds to the Trial object's "config" property.
         """
         with self._engine.connect() as conn:
             cur_params = conn.execute(

From 5cb42580a6627414cd4c7d1f0f6964ab1a18e754 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 18:40:45 +0000
Subject: [PATCH 03/24] comments

---
 mlos_bench/mlos_bench/run.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index a0378e97e7a..e0b5c667f82 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -119,6 +119,8 @@ def _optimize(*,
                 config_id = -1
 
             trial = exp.new_trial(tunables, config={
+                # Add some additional metadata to track for the trial such as the
+                # optimizer config used.
                 "optimizer": opt.name,
                 "opt_target": opt.target,
                 "opt_direction": "min" if opt.is_min else "max",

From a0c21caed1ae0d505a73e5bb474e4070338daaf2 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 18:41:11 +0000
Subject: [PATCH 04/24] data checking

---
 mlos_bench/mlos_bench/storage/base_storage.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlos_bench/mlos_bench/storage/base_storage.py b/mlos_bench/mlos_bench/storage/base_storage.py
index 282d64a429f..b2aec1d13ba 100644
--- a/mlos_bench/mlos_bench/storage/base_storage.py
+++ b/mlos_bench/mlos_bench/storage/base_storage.py
@@ -259,6 +259,7 @@ def __init__(self, *,
             self._trial_id = trial_id
             self._config_id = config_id
             self._opt_target = opt_target
+            assert opt_direction in {"min", "max"}
             self._opt_direction = opt_direction
             self._config = config or {}
 

From cfd3e4e326498db3547f0f4750366bac9161c38a Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 18:41:04 +0000
Subject: [PATCH 05/24] actually do the insert, and add some todo comments

---
 mlos_bench/mlos_bench/optimizers/base_optimizer.py | 2 ++
 mlos_bench/mlos_bench/storage/sql/experiment.py    | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/mlos_bench/mlos_bench/optimizers/base_optimizer.py b/mlos_bench/mlos_bench/optimizers/base_optimizer.py
index 663bef5e7da..ee2e73451f4 100644
--- a/mlos_bench/mlos_bench/optimizers/base_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/base_optimizer.py
@@ -153,6 +153,8 @@ def name(self) -> str:
         """
         return self.__class__.__name__
 
+    # TODO: Expand these properties for multi-objective.
+
     @property
     def is_min(self) -> bool:
         """
diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index 2338161b3b8..0e6b58da80e 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -78,6 +78,11 @@ def _setup(self) -> None:
                     git_commit=self._git_commit,
                     root_env_config=self._root_env_config,
                 ))
+                # TODO: Expand for multiple objectives.
+                conn.execute(self._schema.objectives.insert().values(
+                    objective_target=self._opt_target,
+                    objective_direction=self._opt_direction,
+                ))
             else:
                 if exp_info.trial_id is not None:
                     self._trial_id = exp_info.trial_id + 1

From 7536d59814e743f06054d26716ceceaf31fef95f Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 18:49:29 +0000
Subject: [PATCH 06/24] consolidate logic

---
 mlos_bench/mlos_bench/run.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index e0b5c667f82..e47bc865d68 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -123,7 +123,7 @@ def _optimize(*,
                 # optimizer config used.
                 "optimizer": opt.name,
                 "opt_target": opt.target,
-                "opt_direction": "min" if opt.is_min else "max",
+                "opt_direction": opt.direction,
             })
             _run(env_context, opt_context, trial, global_config)
 

From b14e710e30087f82ff9f0d1365bb5117f774d0e9 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 18:41:29 +0000
Subject: [PATCH 07/24] wip: add a fallback

---
 .../mlos_bench/storage/sql/experiment_data.py | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/mlos_bench/mlos_bench/storage/sql/experiment_data.py b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
index b9b0cc3582b..a4433bd8439 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment_data.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
@@ -34,6 +34,31 @@ def __init__(self, *, engine: Engine, schema: DbSchema, exp_id: str,
         self._engine = engine
         self._schema = schema
 
+    @property
+    def objectives(self) -> Dict[str, str]:
+        if hasattr(self._schema, "objectives"):
+            with self._engine.connect() as conn:
+                objectives = conn.execute(
+                    self._schema.objectives.select().where(
+                        self._schema.objectives.c.exp_id == self._exp_id,
+                    ).order_by(
+                        self._schema.objectives.c.optimization_target.asc(),
+                    )
+                )
+                return {
+                    objective.optimization_target: objective.optimization_direction
+                    for objective in objectives.fetchall()
+                }
+        else:
+            # Backwards compatibility: try and obtain the objectives from the TrialData.
+            # FIXME: convert to metadata
+            return {
+                trial.metadata["opt_target"]: trial.metadata["opt_direction"]
+                for trial in self.trials.values()
+                if trial.metadata.get("opt_target") and trial.metadata.get("opt_direction")
+            }
+
+
     @property
     def trials(self) -> Dict[int, TrialData]:
         with self._engine.connect() as conn:

From 3498ff2bd88d22053e68a362c2ffe5e3dcd3f8b4 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 18:49:51 +0000
Subject: [PATCH 08/24] pylint

---
 mlos_bench/mlos_bench/storage/sql/experiment_data.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/mlos_bench/mlos_bench/storage/sql/experiment_data.py b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
index a4433bd8439..27545541037 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment_data.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
@@ -58,7 +58,6 @@ def objectives(self) -> Dict[str, str]:
                 if trial.metadata.get("opt_target") and trial.metadata.get("opt_direction")
             }
 
-
     @property
     def trials(self) -> Dict[int, TrialData]:
         with self._engine.connect() as conn:

From 4fdf957070bc28e182e020424fe1ce970957366e Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 21:15:52 +0000
Subject: [PATCH 09/24] TODOs

---
 mlos_bench/mlos_bench/run.py                    | 2 ++
 mlos_bench/mlos_bench/storage/sql/experiment.py | 4 ++++
 2 files changed, 6 insertions(+)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index e47bc865d68..001cf223a6b 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -121,6 +121,8 @@ def _optimize(*,
             trial = exp.new_trial(tunables, config={
                 # Add some additional metadata to track for the trial such as the
                 # optimizer config used.
+                # TODO: Improve for supporting multi-objective
+                # (e.g., opt_target_1, opt_target_2, ...)
                 "optimizer": opt.name,
                 "opt_target": opt.target,
                 "opt_direction": opt.direction,
diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index 0e6b58da80e..99e61811591 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -88,6 +88,10 @@ def _setup(self) -> None:
                     self._trial_id = exp_info.trial_id + 1
                 _LOG.info("Continue experiment: %s last trial: %s resume from: %d",
                           self._experiment_id, exp_info.trial_id, self._trial_id)
+                # TODO: Sanity check that certain critical configs (e.g.,
+                # objectives) haven't changed to be incompatible such that a new
+                # experiment should be started (possibly by prewarming with the
+                # previous one).
                 if exp_info.git_commit != self._git_commit:
                     _LOG.warning("Experiment %s git expected: %s %s",
                                  self, exp_info.git_repo, exp_info.git_commit)

From eeea732e5e0dbacc950c8ed3a957239576703cc8 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 21:16:00 +0000
Subject: [PATCH 10/24] fixups

---
 .../mlos_bench/storage/sql/experiment_data.py | 50 +++++++++++++++----
 1 file changed, 39 insertions(+), 11 deletions(-)

diff --git a/mlos_bench/mlos_bench/storage/sql/experiment_data.py b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
index 27545541037..8b75e59d6b4 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment_data.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
@@ -7,6 +7,8 @@
 """
 from typing import Dict
 
+import logging
+
 import pandas
 from sqlalchemy import Engine
 
@@ -16,6 +18,8 @@
 from mlos_bench.storage.base_trial_data import TrialData
 from mlos_bench.storage.sql.trial_data import TrialSqlData
 
+_LOG = logging.getLogger(__name__)
+
 
 class ExperimentSqlData(ExperimentData):
     """
@@ -36,27 +40,51 @@ def __init__(self, *, engine: Engine, schema: DbSchema, exp_id: str,
 
     @property
     def objectives(self) -> Dict[str, str]:
+        objectives: Dict[str, str] = {}
+        # First try to lookup the objectives from the experiment metadata in the storage layer.
         if hasattr(self._schema, "objectives"):
             with self._engine.connect() as conn:
-                objectives = conn.execute(
+                objectives_db_data = conn.execute(
                     self._schema.objectives.select().where(
                         self._schema.objectives.c.exp_id == self._exp_id,
                     ).order_by(
                         self._schema.objectives.c.optimization_target.asc(),
                     )
                 )
-                return {
+                objectives = {
                     objective.optimization_target: objective.optimization_direction
-                    for objective in objectives.fetchall()
+                    for objective in objectives_db_data.fetchall()
                 }
-        else:
-            # Backwards compatibility: try and obtain the objectives from the TrialData.
-            # FIXME: convert to metadata
-            return {
-                trial.metadata["opt_target"]: trial.metadata["opt_direction"]
-                for trial in self.trials.values()
-                if trial.metadata.get("opt_target") and trial.metadata.get("opt_direction")
-            }
+        # Backwards compatibility: try and obtain the objectives from the TrialData and merge them in.
+        # NOTE: The original format of storing opt_target/opt_direction in the Trial
+        # metadata did not support multi-objectives.
+        # Nor does it make it easy to detect when a config change caused a switch in
+        # opt_direction for a given opt_target between run.py executions of an
+        # Experiment.
+        # For now, we simply issue a warning about potentially inconsistent data.
+        for trial in self.trials.values():
+            trial_objectives_df = trial.metadata[
+                trial.metadata["parameter"].isin(("opt_target", "opt_direction"))
+            ][["parameter", "value"]]
+            try:
+                opt_target = trial_objectives_df["opt_target"][0]
+                assert trial_objectives_df["opt_target"].count() == 1, \
+                    "Should only be a single opt_target in the metadata params."
+            except KeyError:
+                # No objective target stored for us to work on, move on.
+                continue
+            try:
+                opt_direction = trial_objectives_df["opt_direction"][0]
+                assert trial_objectives_df["opt_direction"].count() <= 1, \
+                    "Should only be a single opt_direction in the metadata params."
+            except KeyError:
+                pass
+            if opt_target not in objectives:
+                objectives[opt_target] = opt_direction
+            elif opt_direction != objectives[opt_target]:
+                _LOG.warning("Experiment %s has multiple trial optimization directions for optimization_target %s=%s",
+                             self, opt_target, objectives[opt_target])
+        return objectives
 
     @property
     def trials(self) -> Dict[int, TrialData]:

From 6e490102d22f6bf142d76a95d475aab4b249a1f9 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Wed, 10 Jan 2024 21:32:24 +0000
Subject: [PATCH 11/24] add objective info to storage schema

---
 .../mlos_bench/storage/base_experiment_data.py      | 13 +++++++++++++
 mlos_bench/mlos_bench/storage/sql/schema.py         | 13 +++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/mlos_bench/mlos_bench/storage/base_experiment_data.py b/mlos_bench/mlos_bench/storage/base_experiment_data.py
index f9240d8fe7a..29b509fc256 100644
--- a/mlos_bench/mlos_bench/storage/base_experiment_data.py
+++ b/mlos_bench/mlos_bench/storage/base_experiment_data.py
@@ -56,6 +56,19 @@ def root_env_config(self) -> Tuple[str, str, str]:
     def __repr__(self) -> str:
         return f"Experiment :: {self._exp_id}: '{self._description}'"
 
+    @property
+    @abstractmethod
+    def objectives(self) -> Dict[str, str]:
+        """
+        Retrieve the experiment's objectives data from the storage.
+
+        Returns
+        -------
+        objectives : Dict[str, objectives]
+            A dictionary of the experiment's objective names (optimization_targets)
+            and their directions (e.g., min or max).
+        """
+
     @property
     @abstractmethod
     def trials(self) -> Dict[int, TrialData]:
diff --git a/mlos_bench/mlos_bench/storage/sql/schema.py b/mlos_bench/mlos_bench/storage/sql/schema.py
index 7b198e838be..1793272fe8e 100644
--- a/mlos_bench/mlos_bench/storage/sql/schema.py
+++ b/mlos_bench/mlos_bench/storage/sql/schema.py
@@ -62,10 +62,23 @@ def __init__(self, engine: Engine):
             Column("root_env_config", String(1024), nullable=False),
             Column("git_repo", String(1024), nullable=False),
             Column("git_commit", String(40), nullable=False),
+            Column("optimization_target", String(1024), nullable=True),
+            Column("optimization_direction", String(10), nullable=True),
 
             PrimaryKeyConstraint("exp_id"),
         )
 
+        self.objectives = Table(
+            "objectives",
+            self._meta,
+            Column("exp_id"),
+            Column("optimization_target", String(1024), nullable=False),
+            Column("optimization_direction", String(4), nullable=False),
+
+            PrimaryKeyConstraint("exp_id", "optimization_target"),
+            ForeignKeyConstraint(["exp_id"], [self.experiment.c.exp_id]),
+        )
+
         # A workaround for SQLAlchemy issue with autoincrement in DuckDB:
         if engine.dialect.name == "duckdb":
             seq_config_id = Sequence('seq_config_id')

From 687501f0b4f975bfb4392fa5ab394dd6b91594a5 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 18:41:18 +0000
Subject: [PATCH 12/24] todo comments

---
 mlos_bench/mlos_bench/storage/sql/experiment.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index 99e61811591..f13709fa054 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -142,6 +142,8 @@ def load(self, opt_target: Optional[str] = None) -> Tuple[List[dict], List[Optio
                     self._schema.trial.c.trial_id.asc(),
                 )
             )
+            # Note: this iterative approach is somewhat expensive.
+            # TODO: Look into a better bulk fetch option.
             for trial in cur_trials.fetchall():
                 tunables = self._get_params(
                     conn, self._schema.config_param, config_id=trial.config_id)

From aa9e5451f2aaeae43cc6064958dd30092a98beaa Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 21:42:31 +0000
Subject: [PATCH 13/24] stubs for tests

---
 mlos_bench/mlos_bench/tests/storage/exp_data_test.py   | 9 +++++++++
 mlos_bench/mlos_bench/tests/storage/trial_data_test.py | 9 +++++++++
 2 files changed, 18 insertions(+)
 create mode 100644 mlos_bench/mlos_bench/tests/storage/exp_data_test.py
 create mode 100644 mlos_bench/mlos_bench/tests/storage/trial_data_test.py

diff --git a/mlos_bench/mlos_bench/tests/storage/exp_data_test.py b/mlos_bench/mlos_bench/tests/storage/exp_data_test.py
new file mode 100644
index 00000000000..6fbb6f414c0
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/storage/exp_data_test.py
@@ -0,0 +1,9 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Unit tests for loading the experiment metadata.
+"""
+
+# TODO
diff --git a/mlos_bench/mlos_bench/tests/storage/trial_data_test.py b/mlos_bench/mlos_bench/tests/storage/trial_data_test.py
new file mode 100644
index 00000000000..7775e732496
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/storage/trial_data_test.py
@@ -0,0 +1,9 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Unit tests for loading the trial metadata.
+"""
+
+# TODO

From 814a0dd17fb1a11d7fc17db3ee1ce880b842ff03 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 21:54:33 +0000
Subject: [PATCH 14/24] fixup

---
 mlos_bench/mlos_bench/storage/sql/experiment.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index f13709fa054..7e57c96c8e8 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -80,6 +80,7 @@ def _setup(self) -> None:
                 ))
                 # TODO: Expand for multiple objectives.
                 conn.execute(self._schema.objectives.insert().values(
+                    exp_id=self._experiment_id,
                     objective_target=self._opt_target,
                     objective_direction=self._opt_direction,
                 ))

From a2673da7ac95206e6ff1c63faebfb4d0e6072f20 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 22:46:20 +0000
Subject: [PATCH 15/24] move some attrs to the base class

---
 mlos_bench/mlos_bench/storage/base_storage.py | 54 ++++++++++++++++++-
 .../mlos_bench/storage/sql/experiment.py      | 13 +++--
 2 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/mlos_bench/mlos_bench/storage/base_storage.py b/mlos_bench/mlos_bench/storage/base_storage.py
index b2aec1d13ba..3d4fb8db74c 100644
--- a/mlos_bench/mlos_bench/storage/base_storage.py
+++ b/mlos_bench/mlos_bench/storage/base_storage.py
@@ -115,15 +115,28 @@ def experiment(self, *,
         """
 
     class Experiment(metaclass=ABCMeta):
+        # pylint: disable=too-many-instance-attributes
         """
         Base interface for storing the results of the experiment.
         This class is instantiated in the `Storage.experiment()` method.
         """
 
-        def __init__(self, tunables: TunableGroups, experiment_id: str, root_env_config: str):
+        def __init__(self,
+                     *,
+                     tunables: TunableGroups,
+                     experiment_id: str,
+                     trial_id: int,
+                     root_env_config: str,
+                     description: str,
+                     opt_target: str,
+                     opt_direction: str):
             self._tunables = tunables.copy()
+            self._trial_id = trial_id
             self._experiment_id = experiment_id
             (self._git_repo, self._git_commit, self._root_env_config) = get_git_info(root_env_config)
+            self._description = description
+            self._opt_target = opt_target
+            self._opt_direction = opt_direction
 
         def __enter__(self) -> 'Storage.Experiment':
             """
@@ -175,6 +188,31 @@ def _teardown(self, is_ok: bool) -> None:
                 True if there were no exceptions during the experiment, False otherwise.
             """
 
+        @property
+        def experiment_id(self) -> str:
+            """Get the Experiment's ID"""
+            return self._experiment_id
+
+        @property
+        def trial_id(self) -> int:
+            """Get the current Trial ID"""
+            return self._trial_id
+
+        @property
+        def description(self) -> str:
+            """Get the Experiment's description"""
+            return self._description
+
+        @property
+        def opt_target(self) -> str:
+            """Get the Experiment's optimization target"""
+            return self._opt_target
+
+        @property
+        def opt_direction(self) -> str:
+            """Get the Experiment's optimization target"""
+            return self._opt_direction
+
         @abstractmethod
         def merge(self, experiment_ids: List[str]) -> None:
             """
@@ -280,6 +318,20 @@ def config_id(self) -> int:
             """
             return self._config_id
 
+        @property
+        def opt_target(self) -> str:
+            """
+            Get the Trial's optimization target.
+            """
+            return self._opt_target
+
+        @property
+        def opt_direction(self) -> str:
+            """
+            Get the Trial's optimization target.
+            """
+            return self._opt_target
+
         @property
         def tunables(self) -> TunableGroups:
             """
diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index 7e57c96c8e8..b89c09a958f 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -37,13 +37,16 @@ def __init__(self, *,
                  description: str,
                  opt_target: str,
                  opt_direction: str):
-        super().__init__(tunables, experiment_id, root_env_config)
+        super().__init__(
+            tunables=tunables,
+            experiment_id=experiment_id,
+            trial_id=trial_id,
+            root_env_config=root_env_config,
+            description=description,
+            opt_target=opt_target,
+            opt_direction=opt_direction)
         self._engine = engine
         self._schema = schema
-        self._trial_id = trial_id
-        self._description = description
-        self._opt_target = opt_target
-        self._opt_direction = opt_direction
 
     def _setup(self) -> None:
         super()._setup()

From fdf64d2d756a6cc83b1fa4f1782ea98d43da0e9f Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 22:46:35 +0000
Subject: [PATCH 16/24] fixups

---
 mlos_bench/mlos_bench/storage/sql/experiment.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index b89c09a958f..b5307159a29 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -84,8 +84,8 @@ def _setup(self) -> None:
                 # TODO: Expand for multiple objectives.
                 conn.execute(self._schema.objectives.insert().values(
                     exp_id=self._experiment_id,
-                    objective_target=self._opt_target,
-                    objective_direction=self._opt_direction,
+                    optimization_target=self._opt_target,
+                    optimization_direction=self._opt_direction,
                 ))
             else:
                 if exp_info.trial_id is not None:

From d77a8065355f0da4871ed58b903d9a5fd72570f7 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 22:46:43 +0000
Subject: [PATCH 17/24] basic test

---
 mlos_bench/mlos_bench/tests/storage/conftest.py   | 15 ++++++++++++---
 .../mlos_bench/tests/storage/exp_data_test.py     | 13 ++++++++++++-
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/mlos_bench/mlos_bench/tests/storage/conftest.py b/mlos_bench/mlos_bench/tests/storage/conftest.py
index eeb5626f0de..613730cbed7 100644
--- a/mlos_bench/mlos_bench/tests/storage/conftest.py
+++ b/mlos_bench/mlos_bench/tests/storage/conftest.py
@@ -12,13 +12,15 @@
 from mlos_bench.storage.base_storage import Storage
 from mlos_bench.storage.sql.storage import SqlStorage
 
+# pylint: disable=redefined-outer-name
+
 
 @pytest.fixture
-def exp_storage_memory_sql(tunable_groups: TunableGroups) -> Storage.Experiment:
+def storage_memory_sql(tunable_groups: TunableGroups) -> SqlStorage:
     """
     Test fixture for in-memory SQLite3 storage.
     """
-    storage = SqlStorage(
+    return SqlStorage(
         tunables=tunable_groups,
         service=None,
         config={
@@ -26,8 +28,15 @@ def exp_storage_memory_sql(tunable_groups: TunableGroups) -> Storage.Experiment:
             "database": ":memory:",
         }
     )
+
+
+@pytest.fixture
+def exp_storage_memory_sql(storage_memory_sql: Storage) -> SqlStorage.Experiment:
+    """
+    Test fixture for Experiment using in-memory SQLite3 storage.
+    """
     # pylint: disable=unnecessary-dunder-call
-    return storage.experiment(
+    return storage_memory_sql.experiment(
         experiment_id="Test-001",
         trial_id=1,
         root_env_config="environment.jsonc",
diff --git a/mlos_bench/mlos_bench/tests/storage/exp_data_test.py b/mlos_bench/mlos_bench/tests/storage/exp_data_test.py
index 6fbb6f414c0..f6c8eeae680 100644
--- a/mlos_bench/mlos_bench/tests/storage/exp_data_test.py
+++ b/mlos_bench/mlos_bench/tests/storage/exp_data_test.py
@@ -6,4 +6,15 @@
 Unit tests for loading the experiment metadata.
 """
 
-# TODO
+from mlos_bench.storage.base_storage import Storage
+
+
+def test_load_empty_exp_data(storage_memory_sql: Storage, exp_storage_memory_sql: Storage.Experiment) -> None:
+    """
+    Try to retrieve old experimental data from the empty storage.
+    """
+    exp = storage_memory_sql.experiments[exp_storage_memory_sql.experiment_id]
+    assert exp.exp_id == exp_storage_memory_sql.experiment_id
+    assert exp.description == exp_storage_memory_sql.description
+    # Only support single objective for now.
+    assert exp.objectives == {exp_storage_memory_sql.opt_target: exp_storage_memory_sql.opt_direction}

From c15123b4f3d6eef93812102c3a89fd70d4a55985 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 23:34:38 +0000
Subject: [PATCH 18/24] add some more test handling

---
 .../mlos_bench/storage/sql/experiment_data.py | 17 ++++----
 .../tests/storage/trial_data_test.py          | 40 ++++++++++++++++++-
 2 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/mlos_bench/mlos_bench/storage/sql/experiment_data.py b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
index 8b75e59d6b4..f8274ae71f0 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment_data.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
@@ -63,22 +63,23 @@ def objectives(self) -> Dict[str, str]:
         # Experiment.
         # For now, we simply issue a warning about potentially inconsistent data.
         for trial in self.trials.values():
-            trial_objectives_df = trial.metadata[
+            trial_objs_df = trial.metadata[
                 trial.metadata["parameter"].isin(("opt_target", "opt_direction"))
             ][["parameter", "value"]]
             try:
-                opt_target = trial_objectives_df["opt_target"][0]
-                assert trial_objectives_df["opt_target"].count() == 1, \
+                opt_targets = trial_objs_df[trial_objs_df["parameter"] == "opt_target"]
+                assert len(opt_targets) == 1, \
                     "Should only be a single opt_target in the metadata params."
+                opt_target = opt_targets["value"].iloc[0]
             except KeyError:
-                # No objective target stored for us to work on, move on.
                 continue
             try:
-                opt_direction = trial_objectives_df["opt_direction"][0]
-                assert trial_objectives_df["opt_direction"].count() <= 1, \
+                opt_directions = trial_objs_df[trial_objs_df["parameter"] == "opt_direction"]
+                assert len(opt_directions) <= 1, \
                     "Should only be a single opt_direction in the metadata params."
-            except KeyError:
-                pass
+                opt_direction = opt_directions["value"].iloc[0]
+            except (KeyError, IndexError):
+                opt_direction = None
             if opt_target not in objectives:
                 objectives[opt_target] = opt_direction
             elif opt_direction != objectives[opt_target]:
diff --git a/mlos_bench/mlos_bench/tests/storage/trial_data_test.py b/mlos_bench/mlos_bench/tests/storage/trial_data_test.py
index 7775e732496..19f9548e6a2 100644
--- a/mlos_bench/mlos_bench/tests/storage/trial_data_test.py
+++ b/mlos_bench/mlos_bench/tests/storage/trial_data_test.py
@@ -6,4 +6,42 @@
 Unit tests for loading the trial metadata.
 """
 
-# TODO
+from mlos_bench.tunables.tunable_groups import TunableGroups
+from mlos_bench.storage.base_storage import Storage
+
+
+def test_exp_trial_data_objectives(storage_memory_sql: Storage,
+                                   exp_storage_memory_sql: Storage.Experiment,
+                                   tunable_groups: TunableGroups) -> None:
+    """
+    Start a new trial and check the storage for the trial data.
+    """
+
+    trial_opt_new = exp_storage_memory_sql.new_trial(tunable_groups, config={
+        "opt_target": "some-other-target",
+        "opt_direction": "max",
+    })
+    assert trial_opt_new.config() == {
+        "experiment_id": exp_storage_memory_sql.experiment_id,
+        "trial_id": trial_opt_new.trial_id,
+        "opt_target": "some-other-target",
+        "opt_direction": "max",
+    }
+
+    trial_opt_old = exp_storage_memory_sql.new_trial(tunable_groups, config={
+        "opt_target": "back-compat",
+        # "opt_direction": "max",   # missing
+    })
+    assert trial_opt_old.config() == {
+        "experiment_id": exp_storage_memory_sql.experiment_id,
+        "trial_id": trial_opt_old.trial_id,
+        "opt_target": "back-compat",
+    }
+
+    exp = storage_memory_sql.experiments[exp_storage_memory_sql.experiment_id]
+    # objectives should be the combination of both the trial objectives and the experiment objectives
+    assert exp.objectives == {
+        "back-compat": None,
+        "some-other-target": "max",
+        exp_storage_memory_sql.opt_target: exp_storage_memory_sql.opt_direction,
+    }

From 5ffcd670ff885a1a01977662fef08ed63861e176 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 11 Jan 2024 23:40:08 +0000
Subject: [PATCH 19/24] reorg

---
 .../mlos_bench/tests/storage/exp_data_test.py | 38 ++++++++++++++++++
 .../tests/storage/trial_data_test.py          | 40 +------------------
 2 files changed, 39 insertions(+), 39 deletions(-)

diff --git a/mlos_bench/mlos_bench/tests/storage/exp_data_test.py b/mlos_bench/mlos_bench/tests/storage/exp_data_test.py
index f6c8eeae680..fc459172b86 100644
--- a/mlos_bench/mlos_bench/tests/storage/exp_data_test.py
+++ b/mlos_bench/mlos_bench/tests/storage/exp_data_test.py
@@ -7,6 +7,7 @@
 """
 
 from mlos_bench.storage.base_storage import Storage
+from mlos_bench.tunables.tunable_groups import TunableGroups
 
 
 def test_load_empty_exp_data(storage_memory_sql: Storage, exp_storage_memory_sql: Storage.Experiment) -> None:
@@ -18,3 +19,40 @@ def test_load_empty_exp_data(storage_memory_sql: Storage, exp_storage_memory_sql
     assert exp.description == exp_storage_memory_sql.description
     # Only support single objective for now.
     assert exp.objectives == {exp_storage_memory_sql.opt_target: exp_storage_memory_sql.opt_direction}
+
+
+def test_exp_trial_data_objectives(storage_memory_sql: Storage,
+                                   exp_storage_memory_sql: Storage.Experiment,
+                                   tunable_groups: TunableGroups) -> None:
+    """
+    Start a new trial and check the storage for the trial data.
+    """
+
+    trial_opt_new = exp_storage_memory_sql.new_trial(tunable_groups, config={
+        "opt_target": "some-other-target",
+        "opt_direction": "max",
+    })
+    assert trial_opt_new.config() == {
+        "experiment_id": exp_storage_memory_sql.experiment_id,
+        "trial_id": trial_opt_new.trial_id,
+        "opt_target": "some-other-target",
+        "opt_direction": "max",
+    }
+
+    trial_opt_old = exp_storage_memory_sql.new_trial(tunable_groups, config={
+        "opt_target": "back-compat",
+        # "opt_direction": "max",   # missing
+    })
+    assert trial_opt_old.config() == {
+        "experiment_id": exp_storage_memory_sql.experiment_id,
+        "trial_id": trial_opt_old.trial_id,
+        "opt_target": "back-compat",
+    }
+
+    exp = storage_memory_sql.experiments[exp_storage_memory_sql.experiment_id]
+    # objectives should be the combination of both the trial objectives and the experiment objectives
+    assert exp.objectives == {
+        "back-compat": None,
+        "some-other-target": "max",
+        exp_storage_memory_sql.opt_target: exp_storage_memory_sql.opt_direction,
+    }
diff --git a/mlos_bench/mlos_bench/tests/storage/trial_data_test.py b/mlos_bench/mlos_bench/tests/storage/trial_data_test.py
index 19f9548e6a2..7775e732496 100644
--- a/mlos_bench/mlos_bench/tests/storage/trial_data_test.py
+++ b/mlos_bench/mlos_bench/tests/storage/trial_data_test.py
@@ -6,42 +6,4 @@
 Unit tests for loading the trial metadata.
 """
 
-from mlos_bench.tunables.tunable_groups import TunableGroups
-from mlos_bench.storage.base_storage import Storage
-
-
-def test_exp_trial_data_objectives(storage_memory_sql: Storage,
-                                   exp_storage_memory_sql: Storage.Experiment,
-                                   tunable_groups: TunableGroups) -> None:
-    """
-    Start a new trial and check the storage for the trial data.
-    """
-
-    trial_opt_new = exp_storage_memory_sql.new_trial(tunable_groups, config={
-        "opt_target": "some-other-target",
-        "opt_direction": "max",
-    })
-    assert trial_opt_new.config() == {
-        "experiment_id": exp_storage_memory_sql.experiment_id,
-        "trial_id": trial_opt_new.trial_id,
-        "opt_target": "some-other-target",
-        "opt_direction": "max",
-    }
-
-    trial_opt_old = exp_storage_memory_sql.new_trial(tunable_groups, config={
-        "opt_target": "back-compat",
-        # "opt_direction": "max",   # missing
-    })
-    assert trial_opt_old.config() == {
-        "experiment_id": exp_storage_memory_sql.experiment_id,
-        "trial_id": trial_opt_old.trial_id,
-        "opt_target": "back-compat",
-    }
-
-    exp = storage_memory_sql.experiments[exp_storage_memory_sql.experiment_id]
-    # objectives should be the combination of both the trial objectives and the experiment objectives
-    assert exp.objectives == {
-        "back-compat": None,
-        "some-other-target": "max",
-        exp_storage_memory_sql.opt_target: exp_storage_memory_sql.opt_direction,
-    }
+# TODO

From d6ba89e76c1454d8e5373b55cecf683e19cdf579 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@users.noreply.github.com>
Date: Tue, 16 Jan 2024 11:14:30 -0600
Subject: [PATCH 20/24] Update mlos_bench/mlos_bench/run.py

Co-authored-by: Sergiy Matusevych <sergiy.matusevych@gmail.com>
---
 mlos_bench/mlos_bench/run.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index 001cf223a6b..765f9b97bb1 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -122,7 +122,7 @@ def _optimize(*,
                 # Add some additional metadata to track for the trial such as the
                 # optimizer config used.
                 # TODO: Improve for supporting multi-objective
-                # (e.g., opt_target_1, opt_target_2, ...)
+                # (e.g., opt_target_1, opt_target_2, ... and opt_direction_1, opt_direction_2, ...)
                 "optimizer": opt.name,
                 "opt_target": opt.target,
                 "opt_direction": opt.direction,

From 8683aad6e70ce155fc6e1f333191851f7f62ec65 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Tue, 16 Jan 2024 17:20:50 +0000
Subject: [PATCH 21/24] making opt_direction optional

---
 mlos_bench/mlos_bench/storage/base_storage.py | 19 ++++++++++---------
 .../mlos_bench/storage/sql/experiment.py      |  2 +-
 mlos_bench/mlos_bench/storage/sql/storage.py  |  2 +-
 mlos_bench/mlos_bench/storage/sql/trial.py    |  2 +-
 4 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/mlos_bench/mlos_bench/storage/base_storage.py b/mlos_bench/mlos_bench/storage/base_storage.py
index 3d4fb8db74c..1afd699f162 100644
--- a/mlos_bench/mlos_bench/storage/base_storage.py
+++ b/mlos_bench/mlos_bench/storage/base_storage.py
@@ -84,7 +84,7 @@ def experiment(self, *,
                    root_env_config: str,
                    description: str,
                    opt_target: str,
-                   opt_direction: str) -> 'Storage.Experiment':
+                   opt_direction: Optional[str]) -> 'Storage.Experiment':
         """
         Create a new experiment in the storage.
 
@@ -104,7 +104,7 @@ def experiment(self, *,
             Human-readable description of the experiment.
         opt_target : str
             Name of metric we're optimizing for.
-        opt_direction: str
+        opt_direction: Optional[str]
             Direction to optimize the metric (e.g., min or max)
 
         Returns
@@ -129,13 +129,14 @@ def __init__(self,
                      root_env_config: str,
                      description: str,
                      opt_target: str,
-                     opt_direction: str):
+                     opt_direction: Optional[str]):
             self._tunables = tunables.copy()
             self._trial_id = trial_id
             self._experiment_id = experiment_id
             (self._git_repo, self._git_commit, self._root_env_config) = get_git_info(root_env_config)
             self._description = description
             self._opt_target = opt_target
+            assert opt_direction in {None, "min", "max"}
             self._opt_direction = opt_direction
 
         def __enter__(self) -> 'Storage.Experiment':
@@ -209,7 +210,7 @@ def opt_target(self) -> str:
             return self._opt_target
 
         @property
-        def opt_direction(self) -> str:
+        def opt_direction(self) -> Optional[str]:
             """Get the Experiment's optimization target"""
             return self._opt_direction
 
@@ -290,14 +291,14 @@ class Trial(metaclass=ABCMeta):
 
         def __init__(self, *,
                      tunables: TunableGroups, experiment_id: str, trial_id: int,
-                     config_id: int, opt_target: str, opt_direction: str,
+                     config_id: int, opt_target: str, opt_direction: Optional[str],
                      config: Optional[Dict[str, Any]] = None):
             self._tunables = tunables
             self._experiment_id = experiment_id
             self._trial_id = trial_id
             self._config_id = config_id
             self._opt_target = opt_target
-            assert opt_direction in {"min", "max"}
+            assert opt_direction in {None, "min", "max"}
             self._opt_direction = opt_direction
             self._config = config or {}
 
@@ -326,11 +327,11 @@ def opt_target(self) -> str:
             return self._opt_target
 
         @property
-        def opt_direction(self) -> str:
+        def opt_direction(self) -> Optional[str]:
             """
-            Get the Trial's optimization target.
+            Get the Trial's optimization direction (e.g., min or max)
             """
-            return self._opt_target
+            return self._opt_direction
 
         @property
         def tunables(self) -> TunableGroups:
diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index b5307159a29..1bb89e2319e 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -36,7 +36,7 @@ def __init__(self, *,
                  root_env_config: str,
                  description: str,
                  opt_target: str,
-                 opt_direction: str):
+                 opt_direction: Optional[str]):
         super().__init__(
             tunables=tunables,
             experiment_id=experiment_id,
diff --git a/mlos_bench/mlos_bench/storage/sql/storage.py b/mlos_bench/mlos_bench/storage/sql/storage.py
index 64543173c36..9c747f8885c 100644
--- a/mlos_bench/mlos_bench/storage/sql/storage.py
+++ b/mlos_bench/mlos_bench/storage/sql/storage.py
@@ -63,7 +63,7 @@ def experiment(self, *,
                    root_env_config: str,
                    description: str,
                    opt_target: str,
-                   opt_direction: str) -> Storage.Experiment:
+                   opt_direction: Optional[str]) -> Storage.Experiment:
         return Experiment(
             engine=self._engine,
             schema=self._schema,
diff --git a/mlos_bench/mlos_bench/storage/sql/trial.py b/mlos_bench/mlos_bench/storage/sql/trial.py
index 2b6ada313e5..4e13ef8e889 100644
--- a/mlos_bench/mlos_bench/storage/sql/trial.py
+++ b/mlos_bench/mlos_bench/storage/sql/trial.py
@@ -29,7 +29,7 @@ class Trial(Storage.Trial):
     def __init__(self, *,
                  engine: Engine, schema: DbSchema, tunables: TunableGroups,
                  experiment_id: str, trial_id: int, config_id: int,
-                 opt_target: str, opt_direction: str,
+                 opt_target: str, opt_direction: Optional[str],
                  config: Optional[Dict[str, Any]] = None):
         super().__init__(
             tunables=tunables,

From 8d12b72a3b9797d01f6d400bbe46c601c491e45d Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Tue, 16 Jan 2024 17:26:14 +0000
Subject: [PATCH 22/24] adding todo comments and stubbing out weighted
 multi-objective support

---
 mlos_bench/mlos_bench/storage/sql/experiment_data.py | 1 +
 mlos_bench/mlos_bench/storage/sql/schema.py          | 9 +++++++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/mlos_bench/mlos_bench/storage/sql/experiment_data.py b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
index f8274ae71f0..aa90d7fd6b5 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment_data.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment_data.py
@@ -48,6 +48,7 @@ def objectives(self) -> Dict[str, str]:
                     self._schema.objectives.select().where(
                         self._schema.objectives.c.exp_id == self._exp_id,
                     ).order_by(
+                        self._schema.objectives.c.weight.desc(),
                         self._schema.objectives.c.optimization_target.asc(),
                     )
                 )
diff --git a/mlos_bench/mlos_bench/storage/sql/schema.py b/mlos_bench/mlos_bench/storage/sql/schema.py
index 1793272fe8e..504fb83e8aa 100644
--- a/mlos_bench/mlos_bench/storage/sql/schema.py
+++ b/mlos_bench/mlos_bench/storage/sql/schema.py
@@ -11,7 +11,7 @@
 
 from sqlalchemy import (
     Engine, MetaData, Dialect, create_mock_engine,
-    Table, Column, Sequence, Integer, String, DateTime,
+    Table, Column, Sequence, Integer, Float, String, DateTime,
     PrimaryKeyConstraint, ForeignKeyConstraint, UniqueConstraint,
 )
 
@@ -73,7 +73,12 @@ def __init__(self, engine: Engine):
             self._meta,
             Column("exp_id"),
             Column("optimization_target", String(1024), nullable=False),
-            Column("optimization_direction", String(4), nullable=False),
+            Column("optimization_direction", String(4), nullable=True),
+            # TODO: Note: weight is not fully supported yet as currently
+            # multi-objective is expected to explore each objective equally.
+            # Will need to adjust the insert and return values to support this
+            # eventually.
+            Column("weight", Float, nullable=True),
 
             PrimaryKeyConstraint("exp_id", "optimization_target"),
             ForeignKeyConstraint(["exp_id"], [self.experiment.c.exp_id]),

From 667119a5512b6cf37b0fe6d86a35bf7d779cea78 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Tue, 16 Jan 2024 18:00:15 +0000
Subject: [PATCH 23/24] make optimization direction non-nullable

---
 mlos_bench/mlos_bench/storage/sql/schema.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlos_bench/mlos_bench/storage/sql/schema.py b/mlos_bench/mlos_bench/storage/sql/schema.py
index 504fb83e8aa..74d11dabc04 100644
--- a/mlos_bench/mlos_bench/storage/sql/schema.py
+++ b/mlos_bench/mlos_bench/storage/sql/schema.py
@@ -73,7 +73,7 @@ def __init__(self, engine: Engine):
             self._meta,
             Column("exp_id"),
             Column("optimization_target", String(1024), nullable=False),
-            Column("optimization_direction", String(4), nullable=True),
+            Column("optimization_direction", String(4), nullable=False),
             # TODO: Note: weight is not fully supported yet as currently
             # multi-objective is expected to explore each objective equally.
             # Will need to adjust the insert and return values to support this

From d4d8f50cf8a6e899dd212a1e99719811b83513d0 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@users.noreply.github.com>
Date: Tue, 16 Jan 2024 12:01:49 -0600
Subject: [PATCH 24/24] Update
 mlos_bench/mlos_bench/storage/base_experiment_data.py

Co-authored-by: Sergiy Matusevych <sergiy.matusevych@gmail.com>
---
 mlos_bench/mlos_bench/storage/base_experiment_data.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlos_bench/mlos_bench/storage/base_experiment_data.py b/mlos_bench/mlos_bench/storage/base_experiment_data.py
index 29b509fc256..b86347261aa 100644
--- a/mlos_bench/mlos_bench/storage/base_experiment_data.py
+++ b/mlos_bench/mlos_bench/storage/base_experiment_data.py
@@ -64,7 +64,7 @@ def objectives(self) -> Dict[str, str]:
 
         Returns
         -------
-        objectives : Dict[str, objectives]
+        objectives : Dict[str, objective]
             A dictionary of the experiment's objective names (optimization_targets)
             and their directions (e.g., min or max).
         """