From 616f44ef787f2c8fc847133f1a4ca6bc5ba6656e Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Wed, 21 Feb 2024 15:34:05 -0800
Subject: [PATCH 01/18] do not pass the optimizer into _run()

---
 mlos_bench/mlos_bench/run.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index fc78d080554..bc571ad0ab2 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -109,7 +109,8 @@ def _optimize(*,
             opt_context.bulk_register(configs, scores, status)
             # Complete any pending trials.
             for trial in exp.pending_trials(datetime.utcnow(), running=True):
-                _run(env_context, opt_context, trial, global_config)
+                (status, score) = _run(env_context, trial, global_config)
+                opt_context.register(trial.tunables, status, score)
         else:
             _LOG.warning("Skip pending trials and warm-up: %s", opt)
 
@@ -144,7 +145,8 @@ def _optimize(*,
                     "repeat_i": repeat_i,
                     "is_defaults": tunables.is_defaults,
                 })
-                _run(env_context, opt_context, trial, global_config)
+                (status, score) = _run(env_context, trial, global_config)
+                opt_context.register(trial.tunables, status, score)
 
         if do_teardown:
             env_context.teardown()
@@ -154,7 +156,8 @@ def _optimize(*,
     return (best_score, best_config)
 
 
-def _run(env: Environment, opt: Optimizer, trial: Storage.Trial, global_config: Dict[str, Any]) -> None:
+def _run(env: Environment, trial: Storage.Trial,
+         global_config: Dict[str, Any]) -> Tuple[Status, Optional[Dict[str, float]]]:
     """
     Run a single trial.
 
@@ -162,8 +165,6 @@ def _run(env: Environment, opt: Optimizer, trial: Storage.Trial, global_config:
     ----------
     env : Environment
         Benchmarking environment context to run the optimization on.
-    opt : Optimizer
-        An interface to mlos_core optimizers.
     storage : Storage
         A storage system to persist the experiment data.
     global_config : dict
@@ -175,8 +176,7 @@ def _run(env: Environment, opt: Optimizer, trial: Storage.Trial, global_config:
         _LOG.warning("Setup failed: %s :: %s", env, trial.tunables)
         # FIXME: Use the actual timestamp from the environment.
         trial.update(Status.FAILED, datetime.utcnow())
-        opt.register(trial.tunables, Status.FAILED)
-        return
+        return (Status.FAILED, None)
 
     (status, timestamp, results) = env.run()  # Block and wait for the final result.
     _LOG.info("Results: %s :: %s\n%s", trial.tunables, status, results)
@@ -193,7 +193,7 @@ def _run(env: Environment, opt: Optimizer, trial: Storage.Trial, global_config:
     # Filter out non-numeric scores from the optimizer.
     scores = results if not isinstance(results, dict) \
         else {k: float(v) for (k, v) in results.items() if isinstance(v, (int, float))}
-    opt.register(trial.tunables, status, scores)
+    return (status, scores)
 
 
 if __name__ == "__main__":

From 33e332a419ce5ae3d88785a9244441426d01e3ae Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Wed, 21 Feb 2024 15:44:18 -0800
Subject: [PATCH 02/18] mypy fixes

---
 mlos_bench/mlos_bench/run.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index bc571ad0ab2..deb8891ad21 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -109,8 +109,8 @@ def _optimize(*,
             opt_context.bulk_register(configs, scores, status)
             # Complete any pending trials.
             for trial in exp.pending_trials(datetime.utcnow(), running=True):
-                (status, score) = _run(env_context, trial, global_config)
-                opt_context.register(trial.tunables, status, score)
+                (trial_status, trial_score) = _run(env_context, trial, global_config)
+                opt_context.register(trial.tunables, trial_status, trial_score)
         else:
             _LOG.warning("Skip pending trials and warm-up: %s", opt)
 
@@ -145,8 +145,8 @@ def _optimize(*,
                     "repeat_i": repeat_i,
                     "is_defaults": tunables.is_defaults,
                 })
-                (status, score) = _run(env_context, trial, global_config)
-                opt_context.register(trial.tunables, status, score)
+                (trial_status, trial_score) = _run(env_context, trial, global_config)
+                opt_context.register(trial.tunables, trial_status, trial_score)
 
         if do_teardown:
             env_context.teardown()
@@ -169,6 +169,11 @@ def _run(env: Environment, trial: Storage.Trial,
         A storage system to persist the experiment data.
     global_config : dict
         Global configuration parameters.
+
+    Returns
+    -------
+    (trial_status, trial_score) : (Status, Optional[Dict[str, float]])
+        Status and results of the trial.
     """
     _LOG.info("Trial: %s", trial)
 

From 0247259ac964788043d9adfc0bdbb1fe4f8b49ad Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Wed, 21 Feb 2024 18:36:58 -0800
Subject: [PATCH 03/18] start splitting the optimization loop into two

---
 mlos_bench/mlos_bench/run.py | 74 +++++++++++++++++++++++++++++-------
 1 file changed, 60 insertions(+), 14 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index deb8891ad21..49e3eb1ee8c 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -30,7 +30,7 @@ def _main() -> None:
 
     launcher = Launcher("mlos_bench", "Systems autotuning and benchmarking tool")
 
-    result = _optimize(
+    result = _optimization_loop(
         env=launcher.environment,
         opt=launcher.optimizer,
         storage=launcher.storage,
@@ -43,15 +43,15 @@ def _main() -> None:
     _LOG.info("Final result: %s", result)
 
 
-def _optimize(*,
-              env: Environment,
-              opt: Optimizer,
-              storage: Storage,
-              root_env_config: str,
-              global_config: Dict[str, Any],
-              do_teardown: bool,
-              trial_config_repeat_count: int = 1,
-              ) -> Tuple[Optional[float], Optional[TunableGroups]]:
+def _optimization_loop(*,
+                       env: Environment,
+                       opt: Optimizer,
+                       storage: Storage,
+                       root_env_config: str,
+                       global_config: Dict[str, Any],
+                       do_teardown: bool,
+                       trial_config_repeat_count: int = 1,
+                       ) -> Tuple[Optional[float], Optional[TunableGroups]]:
     """
     Main optimization loop.
 
@@ -109,7 +109,7 @@ def _optimize(*,
             opt_context.bulk_register(configs, scores, status)
             # Complete any pending trials.
             for trial in exp.pending_trials(datetime.utcnow(), running=True):
-                (trial_status, trial_score) = _run(env_context, trial, global_config)
+                (trial_status, trial_score) = _run_trial(env_context, trial, global_config)
                 opt_context.register(trial.tunables, trial_status, trial_score)
         else:
             _LOG.warning("Skip pending trials and warm-up: %s", opt)
@@ -145,7 +145,7 @@ def _optimize(*,
                     "repeat_i": repeat_i,
                     "is_defaults": tunables.is_defaults,
                 })
-                (trial_status, trial_score) = _run(env_context, trial, global_config)
+                (trial_status, trial_score) = _run_trial(env_context, trial, global_config)
                 opt_context.register(trial.tunables, trial_status, trial_score)
 
         if do_teardown:
@@ -156,8 +156,54 @@ def _optimize(*,
     return (best_score, best_config)
 
 
-def _run(env: Environment, trial: Storage.Trial,
-         global_config: Dict[str, Any]) -> Tuple[Status, Optional[Dict[str, float]]]:
+def _scheduler(*, exp: Storage.Experiment, env_context: Environment,
+               global_config: Dict[str, Any], running: bool = False) -> None:
+    """
+    Scheduler part of the loop. Check for pending trials in the queue and run them.
+    """
+    for trial in exp.pending_trials(datetime.utcnow(), running=running):
+        _run_trial(env_context, trial, global_config)
+
+
+def _optimizer(exp: Storage.Experiment, opt_context: Optimizer,
+               *, last_trial_id: int = -1, trial_config_repeat_count: int = 1) -> None:
+    """
+    Optimizer part of the loop. Load the results of the executed trials
+    into the optimizer, suggest new configurations, and add them to the queue.
+    """
+    (configs, scores, status) = exp.load(last_trial_id)
+    opt_context.bulk_register(configs, scores, status)
+
+    tunables = opt_context.suggest()
+    _schedule_trial(exp, opt_context, tunables, trial_config_repeat_count)
+
+
+def _schedule_trial(exp: Storage.Experiment, opt: Optimizer,
+                    tunables: TunableGroups, trial_config_repeat_count: int = 1) -> None:
+    """
+    Add one configuration to the queue of trials.
+    """
+    for repeat_i in range(1, trial_config_repeat_count + 1):
+        exp.new_trial(tunables, config={
+            # Add some additional metadata to track for the trial such as the
+            # optimizer config used.
+            # Note: these values are unfortunately mutable at the moment.
+            # Consider them as hints of what the config was the trial *started*.
+            # It is possible that the experiment configs were changed
+            # between resuming the experiment (since that is not currently
+            # prevented).
+            # TODO: Improve for supporting multi-objective
+            # (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,
+            "repeat_i": repeat_i,
+            "is_defaults": tunables.is_defaults,
+        })
+
+
+def _run_trial(env: Environment, trial: Storage.Trial,
+               global_config: Dict[str, Any]) -> Tuple[Status, Optional[Dict[str, float]]]:
     """
     Run a single trial.
 

From 483e378ea8977247c7f18f696acc0ae3e27ab60e Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Thu, 22 Feb 2024 19:18:30 -0800
Subject: [PATCH 04/18] first complete version of the optimization loop (not
 tested yet)

---
 mlos_bench/mlos_bench/run.py | 86 ++++++++++++++++--------------------
 1 file changed, 39 insertions(+), 47 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index 49e3eb1ee8c..bc55b19592d 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -72,7 +72,6 @@ def _optimization_loop(*,
     trial_config_repeat_count : int
         How many trials to repeat for the same configuration.
     """
-    # pylint: disable=too-many-locals
     if trial_config_repeat_count <= 0:
         raise ValueError(f"Invalid trial_config_repeat_count: {trial_config_repeat_count}")
 
@@ -101,52 +100,27 @@ def _optimization_loop(*,
 
         _LOG.info("Experiment: %s Env: %s Optimizer: %s", exp, env, opt)
 
+        last_trial_id = -1
         if opt_context.supports_preload:
-            # Load (tunable values, benchmark scores) to warm-up the optimizer.
-            # `.load()` returns data from ALL merged-in experiments and attempts
-            # to impute the missing tunable values.
-            (configs, scores, status) = exp.load()
-            opt_context.bulk_register(configs, scores, status)
-            # Complete any pending trials.
-            for trial in exp.pending_trials(datetime.utcnow(), running=True):
-                (trial_status, trial_score) = _run_trial(env_context, trial, global_config)
-                opt_context.register(trial.tunables, trial_status, trial_score)
+            # Complete trials that are pending or in-progress.
+            _scheduler(exp, env_context, global_config, running=True)
+            # Load past trials data into the optimizer
+            last_trial_id = _optimizer(exp, opt_context)
         else:
             _LOG.warning("Skip pending trials and warm-up: %s", opt)
 
+        if config_id > 0:
+            tunables = _load_config(exp, env_context, config_id)
+            last_trial_id = _schedule_trial(exp, opt_context, tunables,
+                                            trial_config_repeat_count)
+
         # Now run new trials until the optimizer is done.
         while opt_context.not_converged():
-
-            tunables = opt_context.suggest()
-
-            if config_id > 0:
-                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):
-                    _LOG.debug("Config %d ::\n%s",
-                               config_id, json.dumps(tunable_values, indent=2))
-                config_id = -1
-
-            for repeat_i in range(1, trial_config_repeat_count + 1):
-                trial = exp.new_trial(tunables, config={
-                    # Add some additional metadata to track for the trial such as the
-                    # optimizer config used.
-                    # Note: these values are unfortunately mutable at the moment.
-                    # Consider them as hints of what the config was the trial *started*.
-                    # It is possible that the experiment configs were changed
-                    # between resuming the experiment (since that is not currently
-                    # prevented).
-                    # TODO: Improve for supporting multi-objective
-                    # (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,
-                    "repeat_i": repeat_i,
-                    "is_defaults": tunables.is_defaults,
-                })
-                (trial_status, trial_score) = _run_trial(env_context, trial, global_config)
-                opt_context.register(trial.tunables, trial_status, trial_score)
+            # TODO: In the future, _scheduler and _optimizer
+            # can be run in parallel in two independent loops.
+            _scheduler(exp, env_context, global_config)
+            last_trial_id = _optimizer(exp, opt_context, last_trial_id,
+                                       trial_config_repeat_count)
 
         if do_teardown:
             env_context.teardown()
@@ -156,7 +130,21 @@ def _optimization_loop(*,
     return (best_score, best_config)
 
 
-def _scheduler(*, exp: Storage.Experiment, env_context: Environment,
+def _load_config(exp: Storage.Experiment, env_context: Environment,
+                 config_id: int) -> TunableGroups:
+    """
+    Load the existing tunable configuration from the storage.
+    """
+    tunable_values = exp.load_tunable_config(config_id)
+    tunables = env_context.tunable_params.assign(tunable_values)
+    _LOG.info("Load config from storage: %d", config_id)
+    if _LOG.isEnabledFor(logging.DEBUG):
+        _LOG.debug("Config %d ::\n%s",
+                   config_id, json.dumps(tunable_values, indent=2))
+    return tunables
+
+
+def _scheduler(exp: Storage.Experiment, env_context: Environment,
                global_config: Dict[str, Any], running: bool = False) -> None:
     """
     Scheduler part of the loop. Check for pending trials in the queue and run them.
@@ -166,7 +154,7 @@ def _scheduler(*, exp: Storage.Experiment, env_context: Environment,
 
 
 def _optimizer(exp: Storage.Experiment, opt_context: Optimizer,
-               *, last_trial_id: int = -1, trial_config_repeat_count: int = 1) -> None:
+               last_trial_id: int = -1, trial_config_repeat_count: int = 1) -> int:
     """
     Optimizer part of the loop. Load the results of the executed trials
     into the optimizer, suggest new configurations, and add them to the queue.
@@ -175,16 +163,17 @@ def _optimizer(exp: Storage.Experiment, opt_context: Optimizer,
     opt_context.bulk_register(configs, scores, status)
 
     tunables = opt_context.suggest()
-    _schedule_trial(exp, opt_context, tunables, trial_config_repeat_count)
+    return _schedule_trial(exp, opt_context, tunables, trial_config_repeat_count)
 
 
 def _schedule_trial(exp: Storage.Experiment, opt: Optimizer,
-                    tunables: TunableGroups, trial_config_repeat_count: int = 1) -> None:
+                    tunables: TunableGroups, trial_config_repeat_count: int = 1) -> int:
     """
-    Add one configuration to the queue of trials.
+    Add a configuration to the queue of trials.
     """
+    last_trial_id = -1
     for repeat_i in range(1, trial_config_repeat_count + 1):
-        exp.new_trial(tunables, config={
+        trial = exp.new_trial(tunables, config={
             # Add some additional metadata to track for the trial such as the
             # optimizer config used.
             # Note: these values are unfortunately mutable at the moment.
@@ -200,6 +189,9 @@ def _schedule_trial(exp: Storage.Experiment, opt: Optimizer,
             "repeat_i": repeat_i,
             "is_defaults": tunables.is_defaults,
         })
+        last_trial_id = trial.trial_id
+
+    return last_trial_id
 
 
 def _run_trial(env: Environment, trial: Storage.Trial,

From e97266f149f42c6831e16cf828e9c028967a9a52 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 14:20:58 -0800
Subject: [PATCH 05/18] allow running mlos_bench.run._main directly from unit
 tests + add a unit test for bench (not tested)

---
 mlos_bench/mlos_bench/run.py                   |  3 ++-
 .../mlos_bench/tests/launcher_run_test.py      | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index bc55b19592d..95548a23666 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -26,7 +26,7 @@
 _LOG = logging.getLogger(__name__)
 
 
-def _main() -> None:
+def _main() -> Tuple[Optional[float], Optional[TunableGroups]]:
 
     launcher = Launcher("mlos_bench", "Systems autotuning and benchmarking tool")
 
@@ -41,6 +41,7 @@ def _main() -> None:
     )
 
     _LOG.info("Final result: %s", result)
+    return result
 
 
 def _optimization_loop(*,
diff --git a/mlos_bench/mlos_bench/tests/launcher_run_test.py b/mlos_bench/mlos_bench/tests/launcher_run_test.py
index db8339a6453..1043adb4f42 100644
--- a/mlos_bench/mlos_bench/tests/launcher_run_test.py
+++ b/mlos_bench/mlos_bench/tests/launcher_run_test.py
@@ -10,7 +10,9 @@
 from typing import List
 
 import pytest
+from unittest.mock import MagicMock, patch
 
+from mlos_bench.run import _main
 from mlos_bench.services.local.local_exec import LocalExecService
 from mlos_bench.services.config_persistence import ConfigPersistenceService
 from mlos_bench.util import path_join
@@ -109,3 +111,19 @@ def test_launch_main_app_opt(root_path: str, local_exec_service: LocalExecServic
             r"_optimize INFO Env: Mock environment best score: 64\.53\d+\s*$",
         ]
     )
+
+
+@patch("sys.argv")
+def test_main_bench(mock_argv: MagicMock, root_path: str) -> None:
+    """
+    Run mlos_bench command-line application with given config
+    and check the results in the log.
+    """
+    mock_argv.sys.argv = [
+        "run.py",
+        "--config",
+        "mlos_bench/mlos_bench/tests/config/cli/mock-bench.jsonc",
+    ]
+
+    (score, _config) = _main()
+    assert pytest.approx(score, 1e-6) == 65.67

From 64771fd536f05339cae617de5990046b9813392c Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 14:34:51 -0800
Subject: [PATCH 06/18] move in-process launch to a separate unit test file

---
 .../tests/launcher_in_process_test.py         | 32 +++++++++++++++++++
 .../mlos_bench/tests/launcher_run_test.py     | 18 -----------
 2 files changed, 32 insertions(+), 18 deletions(-)
 create mode 100644 mlos_bench/mlos_bench/tests/launcher_in_process_test.py

diff --git a/mlos_bench/mlos_bench/tests/launcher_in_process_test.py b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
new file mode 100644
index 00000000000..cd88442274a
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
@@ -0,0 +1,32 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Unit tests to check the launcher and the main optimization loop in-process.
+"""
+
+import pytest
+
+from mlos_bench.launcher import Launcher
+from mlos_bench.run import _optimization_loop
+
+
+def test_main_bench() -> None:
+    """
+    Run mlos_bench optimization loop with given config and check the results.
+    """
+    launcher = Launcher("mlos_bench", "TEST RUN", argv=[
+        "--config",
+        "mlos_bench/mlos_bench/tests/config/cli/mock-bench.jsonc",
+    ])
+    (score, _config) = _optimization_loop(
+        env=launcher.environment,
+        opt=launcher.optimizer,
+        storage=launcher.storage,
+        root_env_config=launcher.root_env_config,
+        global_config=launcher.global_config,
+        do_teardown=launcher.teardown,
+        trial_config_repeat_count=launcher.trial_config_repeat_count,
+    )
+    assert pytest.approx(score, 1e-6) == 65.67
diff --git a/mlos_bench/mlos_bench/tests/launcher_run_test.py b/mlos_bench/mlos_bench/tests/launcher_run_test.py
index 1043adb4f42..db8339a6453 100644
--- a/mlos_bench/mlos_bench/tests/launcher_run_test.py
+++ b/mlos_bench/mlos_bench/tests/launcher_run_test.py
@@ -10,9 +10,7 @@
 from typing import List
 
 import pytest
-from unittest.mock import MagicMock, patch
 
-from mlos_bench.run import _main
 from mlos_bench.services.local.local_exec import LocalExecService
 from mlos_bench.services.config_persistence import ConfigPersistenceService
 from mlos_bench.util import path_join
@@ -111,19 +109,3 @@ def test_launch_main_app_opt(root_path: str, local_exec_service: LocalExecServic
             r"_optimize INFO Env: Mock environment best score: 64\.53\d+\s*$",
         ]
     )
-
-
-@patch("sys.argv")
-def test_main_bench(mock_argv: MagicMock, root_path: str) -> None:
-    """
-    Run mlos_bench command-line application with given config
-    and check the results in the log.
-    """
-    mock_argv.sys.argv = [
-        "run.py",
-        "--config",
-        "mlos_bench/mlos_bench/tests/config/cli/mock-bench.jsonc",
-    ]
-
-    (score, _config) = _main()
-    assert pytest.approx(score, 1e-6) == 65.67

From bd7c55e7131f7fd6dca6511463e896108045ea3c Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 15:21:15 -0800
Subject: [PATCH 07/18] add is_warm_up flag to the optimization step

---
 mlos_bench/mlos_bench/optimizers/base_optimizer.py      | 7 +++++--
 mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py | 4 ++--
 mlos_bench/mlos_bench/optimizers/mock_optimizer.py      | 8 +++++---
 mlos_bench/mlos_bench/run.py                            | 7 ++++---
 4 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/mlos_bench/mlos_bench/optimizers/base_optimizer.py b/mlos_bench/mlos_bench/optimizers/base_optimizer.py
index 38d3a0d6c0b..bca4b4f06e2 100644
--- a/mlos_bench/mlos_bench/optimizers/base_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/base_optimizer.py
@@ -195,7 +195,7 @@ def supports_preload(self) -> bool:
 
     @abstractmethod
     def bulk_register(self, configs: Sequence[dict], scores: Sequence[Optional[float]],
-                      status: Optional[Sequence[Status]] = None) -> bool:
+                      status: Optional[Sequence[Status]] = None, is_warm_up: bool = True) -> bool:
         """
         Pre-load the optimizer with the bulk data from previous experiments.
 
@@ -207,13 +207,16 @@ def bulk_register(self, configs: Sequence[dict], scores: Sequence[Optional[float
             Benchmark results from experiments that correspond to `configs`.
         status : Optional[Sequence[float]]
             Status of the experiments that correspond to `configs`.
+        is_warm_up : bool
+            True for the initial load, False for subsequent calls.
 
         Returns
         -------
         is_not_empty : bool
             True if there is data to register, false otherwise.
         """
-        _LOG.info("Warm-up the optimizer with: %d configs, %d scores, %d status values",
+        _LOG.info("%s the optimizer with: %d configs, %d scores, %d status values",
+                  "Warm-up" if is_warm_up else "Load",
                   len(configs or []), len(scores or []), len(status or []))
         if len(configs or []) != len(scores or []):
             raise ValueError("Numbers of configs and scores do not match.")
diff --git a/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py b/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py
index a24745d8f9e..a02b475bac2 100644
--- a/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py
@@ -99,8 +99,8 @@ def name(self) -> str:
         return f"{self.__class__.__name__}:{self._opt.__class__.__name__}"
 
     def bulk_register(self, configs: Sequence[dict], scores: Sequence[Optional[float]],
-                      status: Optional[Sequence[Status]] = None) -> bool:
-        if not super().bulk_register(configs, scores, status):
+                      status: Optional[Sequence[Status]] = None, is_warm_up: bool = True) -> bool:
+        if not super().bulk_register(configs, scores, status, is_warm_up):
             return False
         df_configs = self._to_df(configs)  # Impute missing values, if necessary
         df_scores = pd.Series(scores, dtype=float) * self._opt_sign
diff --git a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py
index 801da49d8fe..11d1b597b17 100644
--- a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py
@@ -42,15 +42,17 @@ def __init__(self,
         self._best_score: Optional[float] = None
 
     def bulk_register(self, configs: Sequence[dict], scores: Sequence[Optional[float]],
-                      status: Optional[Sequence[Status]] = None) -> bool:
-        if not super().bulk_register(configs, scores, status):
+                      status: Optional[Sequence[Status]] = None, is_warm_up: bool = True) -> bool:
+        if not super().bulk_register(configs, scores, status, is_warm_up):
             return False
         if status is None:
             status = [Status.SUCCEEDED] * len(configs)
         for (params, score, trial_status) in zip(configs, scores, status):
             tunables = self._tunables.copy().assign(params)
             self.register(tunables, trial_status, None if score is None else float(score))
-            self._iter -= 1  # Do not advance the iteration counter during warm-up.
+            if is_warm_up:
+                # Do not advance the iteration counter during warm-up.
+                self._iter -= 1
         if _LOG.isEnabledFor(logging.DEBUG):
             (score, _) = self.get_best_observation()
             _LOG.debug("Warm-up end: %s = %s", self.target, score)
diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index 95548a23666..10ba43436cd 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -106,7 +106,7 @@ def _optimization_loop(*,
             # Complete trials that are pending or in-progress.
             _scheduler(exp, env_context, global_config, running=True)
             # Load past trials data into the optimizer
-            last_trial_id = _optimizer(exp, opt_context)
+            last_trial_id = _optimizer(exp, opt_context, is_warm_up=True)
         else:
             _LOG.warning("Skip pending trials and warm-up: %s", opt)
 
@@ -155,13 +155,14 @@ def _scheduler(exp: Storage.Experiment, env_context: Environment,
 
 
 def _optimizer(exp: Storage.Experiment, opt_context: Optimizer,
-               last_trial_id: int = -1, trial_config_repeat_count: int = 1) -> int:
+               last_trial_id: int = -1, trial_config_repeat_count: int = 1,
+               is_warm_up: bool = False) -> int:
     """
     Optimizer part of the loop. Load the results of the executed trials
     into the optimizer, suggest new configurations, and add them to the queue.
     """
     (configs, scores, status) = exp.load(last_trial_id)
-    opt_context.bulk_register(configs, scores, status)
+    opt_context.bulk_register(configs, scores, status, is_warm_up)
 
     tunables = opt_context.suggest()
     return _schedule_trial(exp, opt_context, tunables, trial_config_repeat_count)

From 9f15aeedeb66d2d1ad3a8256186e34b3648b78bc Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 15:37:57 -0800
Subject: [PATCH 08/18] in-process optimizaiton loop invocation works!

---
 mlos_bench/mlos_bench/run.py                            | 2 +-
 mlos_bench/mlos_bench/tests/launcher_in_process_test.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index 10ba43436cd..0f5f0e99700 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -161,7 +161,7 @@ def _optimizer(exp: Storage.Experiment, opt_context: Optimizer,
     Optimizer part of the loop. Load the results of the executed trials
     into the optimizer, suggest new configurations, and add them to the queue.
     """
-    (configs, scores, status) = exp.load(last_trial_id)
+    (configs, scores, status) = exp.load(last_trial_id - 1)
     opt_context.bulk_register(configs, scores, status, is_warm_up)
 
     tunables = opt_context.suggest()
diff --git a/mlos_bench/mlos_bench/tests/launcher_in_process_test.py b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
index cd88442274a..4365a8fc268 100644
--- a/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
+++ b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
@@ -29,4 +29,4 @@ def test_main_bench() -> None:
         do_teardown=launcher.teardown,
         trial_config_repeat_count=launcher.trial_config_repeat_count,
     )
-    assert pytest.approx(score, 1e-6) == 65.67
+    assert pytest.approx(score, 1e-6) == 65.6742

From 65cd07242c74d0df3231a54f4e00061f1221eb32 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 16:03:44 -0800
Subject: [PATCH 09/18] add multi-iteration optimization to in-process test;
 fix the mlos_core bulk registration (check for is_warm_up)

---
 .../optimizers/mlos_core_optimizer.py         |  2 ++
 .../tests/launcher_in_process_test.py         | 23 ++++++++++++++-----
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py b/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py
index a02b475bac2..aee3b7662ff 100644
--- a/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py
@@ -111,6 +111,8 @@ def bulk_register(self, configs: Sequence[dict], scores: Sequence[Optional[float
             df_configs = df_configs[df_status_completed]
             df_scores = df_scores[df_status_completed]
         self._opt.register(df_configs, df_scores)
+        if not is_warm_up:
+            self._iter += len(df_scores)
         if _LOG.isEnabledFor(logging.DEBUG):
             (score, _) = self.get_best_observation()
             _LOG.debug("Warm-up end: %s = %s", self.target, score)
diff --git a/mlos_bench/mlos_bench/tests/launcher_in_process_test.py b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
index 4365a8fc268..fa4979f8182 100644
--- a/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
+++ b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
@@ -6,20 +6,31 @@
 Unit tests to check the launcher and the main optimization loop in-process.
 """
 
+from typing import List
+
 import pytest
 
 from mlos_bench.launcher import Launcher
 from mlos_bench.run import _optimization_loop
 
 
-def test_main_bench() -> None:
+@pytest.mark.parametrize(
+    ("argv", "expected_score"), [
+        ([
+            "--config", "mlos_bench/mlos_bench/tests/config/cli/mock-bench.jsonc",
+        ], 65.6742),
+        ([
+            "--config", "mlos_bench/mlos_bench/tests/config/cli/mock-opt.jsonc",
+            "--trial_config_repeat_count", "3",
+            "--max_iterations", "3",
+        ], 64.53),
+    ]
+)
+def test_main_bench(argv: List[str], expected_score: float) -> None:
     """
     Run mlos_bench optimization loop with given config and check the results.
     """
-    launcher = Launcher("mlos_bench", "TEST RUN", argv=[
-        "--config",
-        "mlos_bench/mlos_bench/tests/config/cli/mock-bench.jsonc",
-    ])
+    launcher = Launcher("mlos_bench", "TEST RUN", argv=argv)
     (score, _config) = _optimization_loop(
         env=launcher.environment,
         opt=launcher.optimizer,
@@ -29,4 +40,4 @@ def test_main_bench() -> None:
         do_teardown=launcher.teardown,
         trial_config_repeat_count=launcher.trial_config_repeat_count,
     )
-    assert pytest.approx(score, 1e-6) == 65.6742
+    assert pytest.approx(score, 1e-6) == expected_score

From c010d957dbc04e3f99d9d81ca19687be36298611 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 16:06:06 -0800
Subject: [PATCH 10/18] make in-process launcerh tests pass

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

diff --git a/mlos_bench/mlos_bench/tests/launcher_in_process_test.py b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
index fa4979f8182..74c9c084d8f 100644
--- a/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
+++ b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
@@ -23,7 +23,7 @@
             "--config", "mlos_bench/mlos_bench/tests/config/cli/mock-opt.jsonc",
             "--trial_config_repeat_count", "3",
             "--max_iterations", "3",
-        ], 64.53),
+        ], 64.8847),
     ]
 )
 def test_main_bench(argv: List[str], expected_score: float) -> None:

From 7cfef3acd7170493b448ad2ac560640048b31e7f Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 16:16:03 -0800
Subject: [PATCH 11/18] remove unnecessary local variables to make pylint happy

---
 mlos_bench/mlos_bench/run.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index 0f5f0e99700..016fd191c94 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -79,10 +79,6 @@ def _optimization_loop(*,
     if _LOG.isEnabledFor(logging.INFO):
         _LOG.info("Root Environment:\n%s", env.pprint())
 
-    experiment_id = global_config["experiment_id"].strip()
-    trial_id = int(global_config["trial_id"])
-    config_id = int(global_config.get("config_id", -1))
-
     # Start new or resume the existing experiment. Verify that the
     # experiment configuration is compatible with the previous runs.
     # If the `merge` config parameter is present, merge in the data
@@ -90,8 +86,8 @@ def _optimization_loop(*,
     with env as env_context, \
          opt as opt_context, \
          storage.experiment(
-            experiment_id=experiment_id,
-            trial_id=trial_id,
+            experiment_id=global_config["experiment_id"].strip(),
+            trial_id=int(global_config["trial_id"]),
             root_env_config=root_env_config,
             description=env.name,
             tunables=env.tunable_params,
@@ -110,6 +106,7 @@ def _optimization_loop(*,
         else:
             _LOG.warning("Skip pending trials and warm-up: %s", opt)
 
+        config_id = int(global_config.get("config_id", -1))
         if config_id > 0:
             tunables = _load_config(exp, env_context, config_id)
             last_trial_id = _schedule_trial(exp, opt_context, tunables,

From 7233180bf50df8d68f3610554dc99f3397936646 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 16:19:19 -0800
Subject: [PATCH 12/18] move trial_config_repeat_count checks to the launcher

---
 mlos_bench/mlos_bench/launcher.py | 6 +++++-
 mlos_bench/mlos_bench/run.py      | 3 ---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/mlos_bench/mlos_bench/launcher.py b/mlos_bench/mlos_bench/launcher.py
index 22dc7d46666..e851581ec80 100644
--- a/mlos_bench/mlos_bench/launcher.py
+++ b/mlos_bench/mlos_bench/launcher.py
@@ -76,7 +76,11 @@ def __init__(self, description: str, long_text: str = "", argv: Optional[List[st
         else:
             config = {}
 
-        self.trial_config_repeat_count: int = args.trial_config_repeat_count or config.get("trial_config_repeat_count", 1)
+        self.trial_config_repeat_count: int = (
+            args.trial_config_repeat_count or config.get("trial_config_repeat_count", 1)
+        )
+        if self.trial_config_repeat_count <= 0:
+            raise ValueError(f"Invalid trial_config_repeat_count: {self.trial_config_repeat_count}")
 
         log_level = args.log_level or config.get("log_level", _LOG_LEVEL)
         try:
diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index 016fd191c94..e6ba50600ae 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -73,9 +73,6 @@ def _optimization_loop(*,
     trial_config_repeat_count : int
         How many trials to repeat for the same configuration.
     """
-    if trial_config_repeat_count <= 0:
-        raise ValueError(f"Invalid trial_config_repeat_count: {trial_config_repeat_count}")
-
     if _LOG.isEnabledFor(logging.INFO):
         _LOG.info("Root Environment:\n%s", env.pprint())
 

From be7dcecd6c7b8be402398a0c26f415b3cd4aa6d6 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 16:45:17 -0800
Subject: [PATCH 13/18] make experiment.load() return trial_ids and use them in
 the optimization loop

---
 mlos_bench/mlos_bench/run.py                  | 21 ++++++++-----------
 mlos_bench/mlos_bench/storage/base_storage.py |  7 ++++---
 .../mlos_bench/storage/sql/experiment.py      |  8 ++++---
 .../mlos_bench/tests/storage/exp_load_test.py |  7 +++++--
 .../tests/storage/trial_schedule_test.py      | 14 +++++++------
 5 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index e6ba50600ae..ed005349f1a 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -99,23 +99,21 @@ def _optimization_loop(*,
             # Complete trials that are pending or in-progress.
             _scheduler(exp, env_context, global_config, running=True)
             # Load past trials data into the optimizer
-            last_trial_id = _optimizer(exp, opt_context, is_warm_up=True)
+            _optimizer(exp, opt_context, is_warm_up=True)
         else:
             _LOG.warning("Skip pending trials and warm-up: %s", opt)
 
         config_id = int(global_config.get("config_id", -1))
         if config_id > 0:
             tunables = _load_config(exp, env_context, config_id)
-            last_trial_id = _schedule_trial(exp, opt_context, tunables,
-                                            trial_config_repeat_count)
+            _schedule_trial(exp, opt_context, tunables, trial_config_repeat_count)
 
         # Now run new trials until the optimizer is done.
         while opt_context.not_converged():
             # TODO: In the future, _scheduler and _optimizer
             # can be run in parallel in two independent loops.
             _scheduler(exp, env_context, global_config)
-            last_trial_id = _optimizer(exp, opt_context, last_trial_id,
-                                       trial_config_repeat_count)
+            _optimizer(exp, opt_context, last_trial_id, trial_config_repeat_count)
 
         if do_teardown:
             env_context.teardown()
@@ -154,20 +152,22 @@ def _optimizer(exp: Storage.Experiment, opt_context: Optimizer,
     """
     Optimizer part of the loop. Load the results of the executed trials
     into the optimizer, suggest new configurations, and add them to the queue.
+    Return the last trial ID processed by the optimizer.
     """
-    (configs, scores, status) = exp.load(last_trial_id - 1)
+    (trial_ids, configs, scores, status) = exp.load(last_trial_id)
     opt_context.bulk_register(configs, scores, status, is_warm_up)
 
     tunables = opt_context.suggest()
-    return _schedule_trial(exp, opt_context, tunables, trial_config_repeat_count)
+    _schedule_trial(exp, opt_context, tunables, trial_config_repeat_count)
+
+    return max(trial_ids, default=last_trial_id)
 
 
 def _schedule_trial(exp: Storage.Experiment, opt: Optimizer,
-                    tunables: TunableGroups, trial_config_repeat_count: int = 1) -> int:
+                    tunables: TunableGroups, trial_config_repeat_count: int = 1) -> None:
     """
     Add a configuration to the queue of trials.
     """
-    last_trial_id = -1
     for repeat_i in range(1, trial_config_repeat_count + 1):
         trial = exp.new_trial(tunables, config={
             # Add some additional metadata to track for the trial such as the
@@ -185,9 +185,6 @@ def _schedule_trial(exp: Storage.Experiment, opt: Optimizer,
             "repeat_i": repeat_i,
             "is_defaults": tunables.is_defaults,
         })
-        last_trial_id = trial.trial_id
-
-    return last_trial_id
 
 
 def _run_trial(env: Environment, trial: Storage.Trial,
diff --git a/mlos_bench/mlos_bench/storage/base_storage.py b/mlos_bench/mlos_bench/storage/base_storage.py
index 687d01ca4ae..e8bc9cdcac8 100644
--- a/mlos_bench/mlos_bench/storage/base_storage.py
+++ b/mlos_bench/mlos_bench/storage/base_storage.py
@@ -258,7 +258,8 @@ def load_telemetry(self, trial_id: int) -> List[Tuple[datetime, str, Any]]:
         @abstractmethod
         def load(self,
                  last_trial_id: int = -1,
-                 opt_target: Optional[str] = None) -> Tuple[List[dict], List[Optional[float]], List[Status]]:
+                 opt_target: Optional[str] = None
+                 ) -> Tuple[List[int], List[dict], List[Optional[float]], List[Status]]:
             """
             Load (tunable values, benchmark scores, status) to warm-up the optimizer.
 
@@ -275,8 +276,8 @@ def load(self,
 
             Returns
             -------
-            (configs, scores, status) : Tuple[List[dict], List[Optional[float]], List[Status]]
-                Tunable values, benchmark scores, and status of the trials.
+            (trial_ids, configs, scores, status) : ([dict], [Optional[float]], [Status])
+                Trial ids, Tunable values, benchmark scores, and status of the trials.
             """
 
         @abstractmethod
diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py
index 7b56a424dca..0be5bc64d27 100644
--- a/mlos_bench/mlos_bench/storage/sql/experiment.py
+++ b/mlos_bench/mlos_bench/storage/sql/experiment.py
@@ -124,9 +124,10 @@ def load_telemetry(self, trial_id: int) -> List[Tuple[datetime, str, Any]]:
 
     def load(self,
              last_trial_id: int = -1,
-             opt_target: Optional[str] = None) -> Tuple[List[dict], List[Optional[float]], List[Status]]:
+             opt_target: Optional[str] = None
+             ) -> Tuple[List[int], List[dict], List[Optional[float]], List[Status]]:
         opt_target = opt_target or self._opt_target
-        (configs, scores, status) = ([], [], [])
+        (trial_ids, configs, scores, status) = ([], [], [], [])
         with self._engine.connect() as conn:
             cur_trials = conn.execute(
                 self._schema.trial.select().with_only_columns(
@@ -154,10 +155,11 @@ def load(self,
             for trial in cur_trials.fetchall():
                 tunables = self._get_params(
                     conn, self._schema.config_param, config_id=trial.config_id)
+                trial_ids.append(trial.trial_id)
                 configs.append(tunables)
                 scores.append(None if trial.metric_value is None else float(trial.metric_value))
                 status.append(Status[trial.status])
-            return (configs, scores, status)
+            return (trial_ids, configs, scores, status)
 
     @staticmethod
     def _get_params(conn: Connection, table: Table, **kwargs: Any) -> Dict[str, Any]:
diff --git a/mlos_bench/mlos_bench/tests/storage/exp_load_test.py b/mlos_bench/mlos_bench/tests/storage/exp_load_test.py
index df19fe77292..3067156a4fe 100644
--- a/mlos_bench/mlos_bench/tests/storage/exp_load_test.py
+++ b/mlos_bench/mlos_bench/tests/storage/exp_load_test.py
@@ -18,7 +18,8 @@ def test_exp_load_empty(exp_storage: Storage.Experiment) -> None:
     """
     Try to retrieve old experimental data from the empty storage.
     """
-    (configs, scores, status) = exp_storage.load()
+    (trial_ids, configs, scores, status) = exp_storage.load()
+    assert not trial_ids
     assert not configs
     assert not scores
     assert not status
@@ -93,6 +94,7 @@ def test_exp_trial_update_categ(exp_storage: Storage.Experiment,
     trial = exp_storage.new_trial(tunable_groups)
     trial.update(Status.SUCCEEDED, datetime.utcnow(), {"score": 99.9, "benchmark": "test"})
     assert exp_storage.load() == (
+        [trial.trial_id],
         [{
             'idle': 'halt',
             'kernel_sched_latency_ns': '2000000',
@@ -133,7 +135,8 @@ def test_exp_trial_pending_3(exp_storage: Storage.Experiment,
     (pending,) = list(exp_storage.pending_trials(datetime.utcnow(), running=True))
     assert pending.trial_id == trial_pend.trial_id
 
-    (configs, scores, status) = exp_storage.load()
+    (trial_ids, configs, scores, status) = exp_storage.load()
+    assert trial_ids == [trial_fail.trial_id, trial_succ.trial_id]
     assert len(configs) == 2
     assert scores == [None, score]
     assert status == [Status.FAILED, Status.SUCCEEDED]
diff --git a/mlos_bench/mlos_bench/tests/storage/trial_schedule_test.py b/mlos_bench/mlos_bench/tests/storage/trial_schedule_test.py
index cac6ddd9b4a..3a582d559d9 100644
--- a/mlos_bench/mlos_bench/tests/storage/trial_schedule_test.py
+++ b/mlos_bench/mlos_bench/tests/storage/trial_schedule_test.py
@@ -72,14 +72,14 @@ def test_schedule_trial(exp_storage: Storage.Experiment,
     # Optimizer side: get trials completed after some known trial:
 
     # No completed trials yet:
-    assert exp_storage.load() == ([], [], [])
+    assert exp_storage.load() == ([], [], [], [])
 
     # Update the status of some trials:
     trial_now1.update(Status.RUNNING, timestamp + timedelta_1min)
     trial_now2.update(Status.RUNNING, timestamp + timedelta_1min)
 
     # Still no completed trials:
-    assert exp_storage.load() == ([], [], [])
+    assert exp_storage.load() == ([], [], [], [])
 
     # Get trials scheduled to run within the next 3 hours:
     pending_ids = _trial_ids(
@@ -107,11 +107,13 @@ def test_schedule_trial(exp_storage: Storage.Experiment,
     trial_1h.update(Status.SUCCEEDED, timestamp + timedelta_1hr * 2, metrics={"score": 1.0})
 
     # Check that three trials have completed so far:
-    (trial_configs, _scores, trial_status) = exp_storage.load()
-    assert len(trial_configs) == 3
+    (trial_ids, trial_configs, trial_scores, trial_status) = exp_storage.load()
+    assert trial_ids == [trial_now1.trial_id, trial_now2.trial_id, trial_1h.trial_id]
+    assert len(trial_configs) == len(trial_scores) == 3
     assert trial_status == [Status.SUCCEEDED, Status.FAILED, Status.SUCCEEDED]
 
     # Get only trials completed after trial_now2:
-    (trial_configs, _scores, trial_status) = exp_storage.load(last_trial_id=trial_now2.trial_id)
-    assert len(trial_configs) == 1
+    (trial_ids, trial_configs, trial_scores, trial_status) = exp_storage.load(last_trial_id=trial_now2.trial_id)
+    assert trial_ids == [trial_1h.trial_id]
+    assert len(trial_configs) == len(trial_scores) == 1
     assert trial_status == [Status.SUCCEEDED]

From 3c52e038a95e70308b1883ce7d84a0f783d9509e Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 16:47:45 -0800
Subject: [PATCH 14/18] use proper last_trial_id in the main loop; fix the unit
 tests

---
 mlos_bench/mlos_bench/run.py                            | 4 ++--
 mlos_bench/mlos_bench/tests/launcher_in_process_test.py | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index ed005349f1a..3293aec3ae5 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -99,7 +99,7 @@ def _optimization_loop(*,
             # Complete trials that are pending or in-progress.
             _scheduler(exp, env_context, global_config, running=True)
             # Load past trials data into the optimizer
-            _optimizer(exp, opt_context, is_warm_up=True)
+            last_trial_id = _optimizer(exp, opt_context, is_warm_up=True)
         else:
             _LOG.warning("Skip pending trials and warm-up: %s", opt)
 
@@ -113,7 +113,7 @@ def _optimization_loop(*,
             # TODO: In the future, _scheduler and _optimizer
             # can be run in parallel in two independent loops.
             _scheduler(exp, env_context, global_config)
-            _optimizer(exp, opt_context, last_trial_id, trial_config_repeat_count)
+            last_trial_id = _optimizer(exp, opt_context, last_trial_id, trial_config_repeat_count)
 
         if do_teardown:
             env_context.teardown()
diff --git a/mlos_bench/mlos_bench/tests/launcher_in_process_test.py b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
index 74c9c084d8f..e8a60ad29cb 100644
--- a/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
+++ b/mlos_bench/mlos_bench/tests/launcher_in_process_test.py
@@ -23,7 +23,7 @@
             "--config", "mlos_bench/mlos_bench/tests/config/cli/mock-opt.jsonc",
             "--trial_config_repeat_count", "3",
             "--max_iterations", "3",
-        ], 64.8847),
+        ], 64.2758),
     ]
 )
 def test_main_bench(argv: List[str], expected_score: float) -> None:

From 0d9dc97aad5f81ec9ce7491cb7563f5637609ad1 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 17:12:59 -0800
Subject: [PATCH 15/18] update launcher tests with the new output patterns

---
 mlos_bench/mlos_bench/tests/launcher_run_test.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/mlos_bench/mlos_bench/tests/launcher_run_test.py b/mlos_bench/mlos_bench/tests/launcher_run_test.py
index db8339a6453..021cead6bab 100644
--- a/mlos_bench/mlos_bench/tests/launcher_run_test.py
+++ b/mlos_bench/mlos_bench/tests/launcher_run_test.py
@@ -81,7 +81,7 @@ def test_launch_main_app_bench(root_path: str, local_exec_service: LocalExecServ
         "--config mlos_bench/mlos_bench/tests/config/cli/mock-bench.jsonc",
         [
             f"^{_RE_DATE} run\\.py:\\d+ " +
-            r"_optimize INFO Env: Mock environment best score: 65\.67\d+\s*$",
+            r"_optimization_loop INFO Env: Mock environment best score: 65\.67\d+\s*$",
         ]
     )
 
@@ -97,15 +97,15 @@ def test_launch_main_app_opt(root_path: str, local_exec_service: LocalExecServic
         [
             # Iteration 1: Expect first value to be the baseline
             f"^{_RE_DATE} mlos_core_optimizer\\.py:\\d+ " +
-            r"register DEBUG Score: 64\.88\d+ Dataframe:\s*$",
+            r"bulk_register DEBUG Warm-up end: score = 64\.88\d+$",
             # Iteration 2: The result may not always be deterministic
             f"^{_RE_DATE} mlos_core_optimizer\\.py:\\d+ " +
-            r"register DEBUG Score: \d+\.\d+ Dataframe:\s*$",
+            r"bulk_register DEBUG Warm-up end: score = \d+\.\d+$",
             # Iteration 3: non-deterministic (depends on the optimizer)
             f"^{_RE_DATE} mlos_core_optimizer\\.py:\\d+ " +
-            r"register DEBUG Score: \d+\.\d+ Dataframe:\s*$",
+            r"bulk_register DEBUG Warm-up end: score = \d+\.\d+$",
             # Final result: baseline is the optimum for the mock environment
             f"^{_RE_DATE} run\\.py:\\d+ " +
-            r"_optimize INFO Env: Mock environment best score: 64\.53\d+\s*$",
+            r"_optimization_loop INFO Env: Mock environment best score: 64\.27\d+\s*$",
         ]
     )

From 4e171e0b5f0f35fc9f22ec3a75a0258b75ef8b69 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 23 Feb 2024 17:21:19 -0800
Subject: [PATCH 16/18] remove unused variable

---
 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 3293aec3ae5..8328a26fdc7 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -169,7 +169,7 @@ def _schedule_trial(exp: Storage.Experiment, opt: Optimizer,
     Add a configuration to the queue of trials.
     """
     for repeat_i in range(1, trial_config_repeat_count + 1):
-        trial = exp.new_trial(tunables, config={
+        exp.new_trial(tunables, config={
             # Add some additional metadata to track for the trial such as the
             # optimizer config used.
             # Note: these values are unfortunately mutable at the moment.

From 52adab86f1a8751d068d4271351ae1a733f72714 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Mon, 26 Feb 2024 13:43:34 -0800
Subject: [PATCH 17/18] better naming for functions in the optimization loop

---
 mlos_bench/mlos_bench/run.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py
index 8328a26fdc7..18bf779180e 100755
--- a/mlos_bench/mlos_bench/run.py
+++ b/mlos_bench/mlos_bench/run.py
@@ -97,9 +97,9 @@ def _optimization_loop(*,
         last_trial_id = -1
         if opt_context.supports_preload:
             # Complete trials that are pending or in-progress.
-            _scheduler(exp, env_context, global_config, running=True)
+            _run_schedule(exp, env_context, global_config, running=True)
             # Load past trials data into the optimizer
-            last_trial_id = _optimizer(exp, opt_context, is_warm_up=True)
+            last_trial_id = _get_optimizer_suggestions(exp, opt_context, is_warm_up=True)
         else:
             _LOG.warning("Skip pending trials and warm-up: %s", opt)
 
@@ -112,8 +112,8 @@ def _optimization_loop(*,
         while opt_context.not_converged():
             # TODO: In the future, _scheduler and _optimizer
             # can be run in parallel in two independent loops.
-            _scheduler(exp, env_context, global_config)
-            last_trial_id = _optimizer(exp, opt_context, last_trial_id, trial_config_repeat_count)
+            _run_schedule(exp, env_context, global_config)
+            last_trial_id = _get_optimizer_suggestions(exp, opt_context, last_trial_id, trial_config_repeat_count)
 
         if do_teardown:
             env_context.teardown()
@@ -137,8 +137,8 @@ def _load_config(exp: Storage.Experiment, env_context: Environment,
     return tunables
 
 
-def _scheduler(exp: Storage.Experiment, env_context: Environment,
-               global_config: Dict[str, Any], running: bool = False) -> None:
+def _run_schedule(exp: Storage.Experiment, env_context: Environment,
+                  global_config: Dict[str, Any], running: bool = False) -> None:
     """
     Scheduler part of the loop. Check for pending trials in the queue and run them.
     """
@@ -146,9 +146,9 @@ def _scheduler(exp: Storage.Experiment, env_context: Environment,
         _run_trial(env_context, trial, global_config)
 
 
-def _optimizer(exp: Storage.Experiment, opt_context: Optimizer,
-               last_trial_id: int = -1, trial_config_repeat_count: int = 1,
-               is_warm_up: bool = False) -> int:
+def _get_optimizer_suggestions(exp: Storage.Experiment, opt_context: Optimizer,
+                               last_trial_id: int = -1, trial_config_repeat_count: int = 1,
+                               is_warm_up: bool = False) -> int:
     """
     Optimizer part of the loop. Load the results of the executed trials
     into the optimizer, suggest new configurations, and add them to the queue.

From 5aca76450973458c35102bf6cbce1ca66cc0ab6b Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Mon, 26 Feb 2024 16:10:59 -0800
Subject: [PATCH 18/18] change the default value for is_warm_up parameter to
 False

---
 mlos_bench/mlos_bench/optimizers/base_optimizer.py      | 2 +-
 mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py | 2 +-
 mlos_bench/mlos_bench/optimizers/mock_optimizer.py      | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlos_bench/mlos_bench/optimizers/base_optimizer.py b/mlos_bench/mlos_bench/optimizers/base_optimizer.py
index bca4b4f06e2..e41693e88ab 100644
--- a/mlos_bench/mlos_bench/optimizers/base_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/base_optimizer.py
@@ -195,7 +195,7 @@ def supports_preload(self) -> bool:
 
     @abstractmethod
     def bulk_register(self, configs: Sequence[dict], scores: Sequence[Optional[float]],
-                      status: Optional[Sequence[Status]] = None, is_warm_up: bool = True) -> bool:
+                      status: Optional[Sequence[Status]] = None, is_warm_up: bool = False) -> bool:
         """
         Pre-load the optimizer with the bulk data from previous experiments.
 
diff --git a/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py b/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py
index aee3b7662ff..6c1f8854446 100644
--- a/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/mlos_core_optimizer.py
@@ -99,7 +99,7 @@ def name(self) -> str:
         return f"{self.__class__.__name__}:{self._opt.__class__.__name__}"
 
     def bulk_register(self, configs: Sequence[dict], scores: Sequence[Optional[float]],
-                      status: Optional[Sequence[Status]] = None, is_warm_up: bool = True) -> bool:
+                      status: Optional[Sequence[Status]] = None, is_warm_up: bool = False) -> bool:
         if not super().bulk_register(configs, scores, status, is_warm_up):
             return False
         df_configs = self._to_df(configs)  # Impute missing values, if necessary
diff --git a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py
index 11d1b597b17..ed156fe44b8 100644
--- a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py
@@ -42,7 +42,7 @@ def __init__(self,
         self._best_score: Optional[float] = None
 
     def bulk_register(self, configs: Sequence[dict], scores: Sequence[Optional[float]],
-                      status: Optional[Sequence[Status]] = None, is_warm_up: bool = True) -> bool:
+                      status: Optional[Sequence[Status]] = None, is_warm_up: bool = False) -> bool:
         if not super().bulk_register(configs, scores, status, is_warm_up):
             return False
         if status is None: