Repository: microsoft/MLOS

Issue:
mlos_bench: grid search support
In benchmarking only mode, sometimes we actually do want to use grid search to evaluate all possibilities for comparison sake.

This can also be useful when comparing more efficient optimizer search strategy performance as well.

We should be able to implement this as an additional mlos_bench optimizer in combination with the storage backend to check which ones have already been evaluated.


Discussion:
This btw would be a great place for parallel trial execution (#380)

Resolving patch:
diff --git a/mlos_bench/mlos_bench/optimizers/grid_search_optimizer.py b/mlos_bench/mlos_bench/optimizers/grid_search_optimizer.py
new file mode 100644
--- /dev/null
+++ b/mlos_bench/mlos_bench/optimizers/grid_search_optimizer.py
@@ -0,0 +1,188 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Grid search optimizer for mlos_bench.
+"""
+
+import logging
+
+from typing import Dict, Iterable, Set, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import ConfigSpace
+from ConfigSpace.util import generate_grid
+
+from mlos_bench.environments.status import Status
+from mlos_bench.tunables.tunable import TunableValue
+from mlos_bench.tunables.tunable_groups import TunableGroups
+from mlos_bench.optimizers.base_optimizer import Optimizer
+from mlos_bench.optimizers.convert_configspace import configspace_data_to_tunable_values
+from mlos_bench.services.base_service import Service
+
+_LOG = logging.getLogger(__name__)
+
+
+class GridSearchOptimizer(Optimizer):
+    """
+    Grid search optimizer.
+    """
+
+    def __init__(self,
+                 tunables: TunableGroups,
+                 config: dict,
+                 global_config: Optional[dict] = None,
+                 service: Optional[Service] = None):
+        super().__init__(tunables, config, global_config, service)
+
+        self._best_config: Optional[TunableGroups] = None
+        self._best_score: Optional[float] = None
+
+        # Track the grid as a set of tuples of tunable values and reconstruct the
+        # dicts as necessary.
+        # Note: this is not the most effecient way to do this, but avoids
+        # introducing a new data structure for hashable dicts.
+        # See https://github.com/microsoft/MLOS/pull/690 for further discussion.
+
+        self._sanity_check()
+        # The ordered set of pending configs that have not yet been suggested.
+        self._config_keys, self._pending_configs = self._get_grid()
+        assert self._pending_configs
+        # A set of suggested configs that have not yet been registered.
+        self._suggested_configs: Set[Tuple[TunableValue, ...]] = set()
+
+    def _sanity_check(self) -> None:
+        size = np.prod([tunable.cardinality for (tunable, _group) in self._tunables])
+        if size == np.inf:
+            raise ValueError(f"Unquantized tunables are not supported for grid search: {self._tunables}")
+        if size > 10000:
+            _LOG.warning("Large number %d of config points requested for grid search: %s", size, self._tunables)
+        if size > self._max_iter:
+            _LOG.warning("Grid search size %d, is greater than max iterations %d", size, self._max_iter)
+
+    def _get_grid(self) -> Tuple[Tuple[str, ...], Dict[Tuple[TunableValue, ...], None]]:
+        """
+        Gets a grid of configs to try.
+
+        Order is given by ConfigSpace, but preserved by dict ordering semantics.
+        """
+        # Since we are using ConfigSpace to generate the grid, but only tracking the
+        # values as (ordered) tuples, we also need to use its ordering on column
+        # names instead of the order given by TunableGroups.
+        configs = [
+            configspace_data_to_tunable_values(dict(config))
+            for config in
+            generate_grid(self.config_space, {
+                tunable.name: int(tunable.cardinality)
+                for (tunable, _group) in self._tunables
+                if tunable.quantization or tunable.type == "int"
+            })
+        ]
+        names = set(tuple(configs.keys()) for configs in configs)
+        assert len(names) == 1
+        return names.pop(), {tuple(configs.values()): None for configs in configs}
+
+    @property
+    def pending_configs(self) -> Iterable[Dict[str, TunableValue]]:
+        """
+        Gets the set of pending configs in this grid search optimizer.
+
+        Returns
+        -------
+        Iterable[Dict[str, TunableValue]]
+        """
+        # See NOTEs above.
+        return (dict(zip(self._config_keys, config)) for config in self._pending_configs.keys())
+
+    @property
+    def suggested_configs(self) -> Iterable[Dict[str, TunableValue]]:
+        """
+        Gets the set of configs that have been suggested but not yet registered.
+
+        Returns
+        -------
+        Iterable[Dict[str, TunableValue]]
+        """
+        # See NOTEs above.
+        return (dict(zip(self._config_keys, config)) for config in self._suggested_configs)
+
+    def bulk_register(self, configs: Sequence[dict], scores: Sequence[Optional[float]],
+                      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:
+            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))
+            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("%s end: %s = %s", "Warm-up" if is_warm_up else "Update", self.target, score)
+        return True
+
+    def suggest(self) -> TunableGroups:
+        """
+        Generate the next grid search suggestion.
+        """
+        tunables = self._tunables.copy()
+        if self._start_with_defaults:
+            _LOG.info("Use default values for the first trial")
+            self._start_with_defaults = False
+            tunables = tunables.restore_defaults()
+            # Need to index based on ConfigSpace dict ordering.
+            default_config = dict(self.config_space.get_default_configuration())
+            assert tunables.get_param_values() == default_config
+            # Move the default from the pending to the suggested set.
+            default_config_values = tuple(default_config.values())
+            del self._pending_configs[default_config_values]
+            self._suggested_configs.add(default_config_values)
+        else:
+            # Select the first item from the pending configs.
+            if not self._pending_configs and self._iter <= self._max_iter:
+                _LOG.info("No more pending configs to suggest. Restarting grid.")
+                self._config_keys, self._pending_configs = self._get_grid()
+            try:
+                next_config_values = next(iter(self._pending_configs.keys()))
+            except StopIteration as exc:
+                raise ValueError("No more pending configs to suggest.") from exc
+            next_config = dict(zip(self._config_keys, next_config_values))
+            tunables.assign(next_config)
+            # Move it to the suggested set.
+            self._suggested_configs.add(next_config_values)
+            del self._pending_configs[next_config_values]
+        _LOG.info("Iteration %d :: Suggest: %s", self._iter, tunables)
+        return tunables
+
+    def register(self, tunables: TunableGroups, status: Status,
+                 score: Optional[Union[float, dict]] = None) -> Optional[float]:
+        registered_score = super().register(tunables, status, score)
+        if status.is_succeeded() and (
+            self._best_score is None or (registered_score is not None and registered_score < self._best_score)
+        ):
+            self._best_score = registered_score
+            self._best_config = tunables.copy()
+        self._iter += 1
+        try:
+            config = dict(ConfigSpace.Configuration(self.config_space, values=tunables.get_param_values()))
+            self._suggested_configs.remove(tuple(config.values()))
+        except KeyError:
+            _LOG.warning("Attempted to remove missing config (previously registered?) from suggested set: %s", tunables)
+        return registered_score
+
+    def not_converged(self) -> bool:
+        if self._iter > self._max_iter:
+            if bool(self._pending_configs):
+                _LOG.warning("Exceeded max iterations, but still have %d pending configs: %s",
+                             len(self._pending_configs), list(self._pending_configs.keys()))
+            return False
+        return bool(self._pending_configs)
+
+    def get_best_observation(self) -> Union[Tuple[float, TunableGroups], Tuple[None, None]]:
+        if self._best_score is None:
+            return (None, None)
+        assert self._best_config is not None
+        return (self._best_score * self._opt_sign, self._best_config)
diff --git a/mlos_bench/mlos_bench/tunables/tunable_groups.py b/mlos_bench/mlos_bench/tunables/tunable_groups.py
--- a/mlos_bench/mlos_bench/tunables/tunable_groups.py
+++ b/mlos_bench/mlos_bench/tunables/tunable_groups.py
@@ -39,6 +39,9 @@ def __init__(self, config: Optional[dict] = None):
     def __bool__(self) -> bool:
         return bool(self._index)
 
+    def __len__(self) -> int:
+        return len(self._index)
+
     def __eq__(self, other: object) -> bool:
         """
         Check if two TunableGroups are equal.


Test patch:
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/optimizers/test-cases/bad/invalid/grid_search_opt_invalid_max_iterations.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/optimizers/test-cases/bad/invalid/grid_search_opt_invalid_max_iterations.jsonc
new file mode 100644
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/config/schemas/optimizers/test-cases/bad/invalid/grid_search_opt_invalid_max_iterations.jsonc
@@ -0,0 +1,6 @@
+{
+    "class": "mlos_bench.optimizers.grid_search_optimizer.GridSearchOptimizer",
+    "config": {
+        "max_iterations": null,
+    }
+}
\ No newline at end of file
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/optimizers/test-cases/bad/unhandled/grid_search_opt_extra.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/optimizers/test-cases/bad/unhandled/grid_search_opt_extra.jsonc
new file mode 100644
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/config/schemas/optimizers/test-cases/bad/unhandled/grid_search_opt_extra.jsonc
@@ -0,0 +1,6 @@
+{
+    "class": "mlos_bench.optimizers.grid_search_optimizer.GridSearchOptimizer",
+    "config": {
+        "extra": "unhandled"
+    }
+}
\ No newline at end of file
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/optimizers/test-cases/good/full/grid_search_opt_full.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/optimizers/test-cases/good/full/grid_search_opt_full.jsonc
new file mode 100644
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/config/schemas/optimizers/test-cases/good/full/grid_search_opt_full.jsonc
@@ -0,0 +1,14 @@
+{
+    "class": "mlos_bench.optimizers.grid_search_optimizer.GridSearchOptimizer",
+    "description": "GridSearchOptimizer",
+    "include_tunables": [
+        "some/path/to/tunables.jsonc"
+    ],
+    "config": {
+        "max_iterations": 100,
+        "optimization_direction": "max",
+        "optimization_target": "score",
+        "seed": 12345,
+        "start_with_defaults": true
+    }
+}
\ No newline at end of file
diff --git a/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py b/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py
new file mode 100644
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py
@@ -0,0 +1,250 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Unit tests for grid search mlos_bench optimizer.
+"""
+
+from typing import Dict, List
+
+import itertools
+import math
+import random
+
+import pytest
+
+from mlos_bench.environments.status import Status
+from mlos_bench.optimizers.grid_search_optimizer import GridSearchOptimizer
+from mlos_bench.tunables.tunable import TunableValue
+from mlos_bench.tunables.tunable_groups import TunableGroups
+
+
+# pylint: disable=redefined-outer-name
+
+@pytest.fixture
+def grid_search_tunables_config() -> dict:
+    """
+    Test fixture for grid search optimizer tunables config.
+    """
+    return {
+        "grid": {
+            "cost": 1,
+            "params": {
+                "cat": {
+                    "type": "categorical",
+                    "values": ["a", "b", "c"],
+                    "default": "a",
+                },
+                "int": {
+                    "type": "int",
+                    "range": [1, 3],
+                    "default": 2,
+                },
+                "float": {
+                    "type": "float",
+                    "range": [0, 1],
+                    "default": 0.5,
+                    "quantization": 0.25,
+                },
+            },
+        },
+    }
+
+
+@pytest.fixture
+def grid_search_tunables_grid(grid_search_tunables: TunableGroups) -> List[Dict[str, TunableValue]]:
+    """
+    Test fixture for grid from tunable groups.
+    Used to check that the grids are the same (ignoring order).
+    """
+    tunables_params_values = [tunable.values for tunable, _group in grid_search_tunables if tunable.values is not None]
+    tunable_names = tuple(tunable.name for tunable, _group in grid_search_tunables if tunable.values is not None)
+    return list(dict(zip(tunable_names, combo)) for combo in itertools.product(*tunables_params_values))
+
+
+@pytest.fixture
+def grid_search_tunables(grid_search_tunables_config: dict) -> TunableGroups:
+    """
+    Test fixture for grid search optimizer tunables.
+    """
+    return TunableGroups(grid_search_tunables_config)
+
+
+@pytest.fixture
+def grid_search_opt(grid_search_tunables: TunableGroups,
+                    grid_search_tunables_grid: List[Dict[str, TunableValue]]) -> GridSearchOptimizer:
+    """
+    Test fixture for grid search optimizer.
+    """
+    assert len(grid_search_tunables) == 3
+    # Test the convergence logic by controlling the number of iterations to be not a
+    # multiple of the number of elements in the grid.
+    max_iterations = len(grid_search_tunables_grid) * 2 - 3
+    return GridSearchOptimizer(tunables=grid_search_tunables, config={
+        "max_iterations": max_iterations,
+        "optimization_direction": "max",
+    })
+
+
+def test_grid_search_grid(grid_search_opt: GridSearchOptimizer,
+                          grid_search_tunables: TunableGroups,
+                          grid_search_tunables_grid: List[Dict[str, TunableValue]]) -> None:
+    """
+    Make sure that grid search optimizer initializes and works correctly.
+    """
+    # Check the size.
+    expected_grid_size = math.prod(tunable.cardinality for tunable, _group in grid_search_tunables)
+    assert expected_grid_size > len(grid_search_tunables)
+    assert len(grid_search_tunables_grid) == expected_grid_size
+    # Check for specific example configs inclusion.
+    expected_config_example: Dict[str, TunableValue] = {
+        "cat": "a",
+        "int": 2,
+        "float": 0.75,
+    }
+    grid_search_opt_pending_configs = list(grid_search_opt.pending_configs)
+    assert expected_config_example in grid_search_tunables_grid
+    assert expected_config_example in grid_search_opt_pending_configs
+    # Check the rest of the contents.
+    # Note: ConfigSpace param name vs TunableGroup parameter name order is not
+    # consistent, so we need to full dict comparison.
+    assert len(grid_search_opt_pending_configs) == expected_grid_size
+    assert all(config in grid_search_tunables_grid for config in grid_search_opt_pending_configs)
+    assert all(config in grid_search_opt_pending_configs for config in grid_search_tunables_grid)
+    # Order is less relevant to us, so we'll just check that the sets are the same.
+    # assert grid_search_opt.pending_configs == grid_search_tunables_grid
+
+
+def test_grid_search(grid_search_opt: GridSearchOptimizer,
+                     grid_search_tunables: TunableGroups,
+                     grid_search_tunables_grid: List[Dict[str, TunableValue]]) -> None:
+    """
+    Make sure that grid search optimizer initializes and works correctly.
+    """
+    score = 1.0
+    status = Status.SUCCEEDED
+    suggestion = grid_search_opt.suggest()
+    suggestion_dict = suggestion.get_param_values()
+    default_config = grid_search_tunables.restore_defaults().get_param_values()
+
+    # First suggestion should be the defaults.
+    assert suggestion.get_param_values() == default_config
+    # But that shouldn't be the first element in the grid search.
+    assert suggestion_dict != next(iter(grid_search_tunables_grid))
+    # The suggestion should no longer be in the pending_configs.
+    assert suggestion_dict not in grid_search_opt.pending_configs
+    # But it should be in the suggested_configs now (and the only one).
+    assert list(grid_search_opt.suggested_configs) == [default_config]
+
+    # Register a score for that suggestion.
+    grid_search_opt.register(suggestion, status, score)
+    # Now it shouldn't be in the suggested_configs.
+    assert len(list(grid_search_opt.suggested_configs)) == 0
+
+    grid_search_tunables_grid.remove(default_config)
+    assert default_config not in grid_search_opt.pending_configs
+    assert all(config in grid_search_tunables_grid for config in grid_search_opt.pending_configs)
+    assert all(config in list(grid_search_opt.pending_configs) for config in grid_search_tunables_grid)
+
+    # The next suggestion should be a different element in the grid search.
+    suggestion = grid_search_opt.suggest()
+    suggestion_dict = suggestion.get_param_values()
+    assert suggestion_dict != default_config
+    assert suggestion_dict not in grid_search_opt.pending_configs
+    assert suggestion_dict in grid_search_opt.suggested_configs
+    grid_search_opt.register(suggestion, status, score)
+    assert suggestion_dict not in grid_search_opt.pending_configs
+    assert suggestion_dict not in grid_search_opt.suggested_configs
+
+    grid_search_tunables_grid.remove(suggestion.get_param_values())
+    assert all(config in grid_search_tunables_grid for config in grid_search_opt.pending_configs)
+    assert all(config in list(grid_search_opt.pending_configs) for config in grid_search_tunables_grid)
+
+    # FIXME: Should we consider not_converged as the "max_iterations", an empty grid, or both?
+
+    # Try to empty the rest of the grid.
+    while grid_search_opt.not_converged():
+        suggestion = grid_search_opt.suggest()
+        grid_search_opt.register(suggestion, status, score)
+
+    # The grid search should be empty now.
+    assert not list(grid_search_opt.pending_configs)
+    assert not list(grid_search_opt.suggested_configs)
+    assert not grid_search_opt.not_converged()
+
+    # But if we still have iterations left, we should be able to suggest again by refilling the grid.
+    assert grid_search_opt.current_iteration < grid_search_opt.max_iterations
+    assert grid_search_opt.suggest()
+    assert list(grid_search_opt.pending_configs)
+    assert list(grid_search_opt.suggested_configs)
+    assert grid_search_opt.not_converged()
+
+    # Try to finish the rest of our iterations by repeating the grid.
+    while grid_search_opt.not_converged():
+        suggestion = grid_search_opt.suggest()
+        grid_search_opt.register(suggestion, status, score)
+    assert not grid_search_opt.not_converged()
+    assert grid_search_opt.current_iteration >= grid_search_opt.max_iterations
+    assert list(grid_search_opt.pending_configs)
+    assert list(grid_search_opt.suggested_configs)
+
+
+def test_grid_search_async_order(grid_search_opt: GridSearchOptimizer) -> None:
+    """
+    Make sure that grid search optimizer works correctly when suggest and register
+    are called out of order.
+    """
+    score = 1.0
+    status = Status.SUCCEEDED
+    suggest_count = 10
+    suggested = [grid_search_opt.suggest() for _ in range(suggest_count)]
+    suggested_shuffled = suggested.copy()
+    # Try to ensure the shuffled list is different.
+    for _ in range(3):
+        random.shuffle(suggested_shuffled)
+        if suggested_shuffled != suggested:
+            break
+    assert suggested != suggested_shuffled
+
+    for suggestion in suggested_shuffled:
+        suggestion_dict = suggestion.get_param_values()
+        assert suggestion_dict not in grid_search_opt.pending_configs
+        assert suggestion_dict in grid_search_opt.suggested_configs
+        grid_search_opt.register(suggestion, status, score)
+        assert suggestion_dict not in grid_search_opt.suggested_configs
+
+    best_score, best_config = grid_search_opt.get_best_observation()
+    assert best_score == score
+
+    # test re-register with higher score
+    best_suggestion = suggested_shuffled[0]
+    best_suggestion_dict = best_suggestion.get_param_values()
+    assert best_suggestion_dict not in grid_search_opt.pending_configs
+    assert best_suggestion_dict not in grid_search_opt.suggested_configs
+    best_suggestion_score = score - 1 if grid_search_opt.direction == "min" else score + 1
+    grid_search_opt.register(best_suggestion, status, best_suggestion_score)
+    assert best_suggestion_dict not in grid_search_opt.suggested_configs
+
+    best_score, best_config = grid_search_opt.get_best_observation()
+    assert best_score == best_suggestion_score
+    assert best_config == best_suggestion
+
+    # Check bulk register
+    suggested = [grid_search_opt.suggest() for _ in range(suggest_count)]
+    assert all(suggestion.get_param_values() not in grid_search_opt.pending_configs for suggestion in suggested)
+    assert all(suggestion.get_param_values() in grid_search_opt.suggested_configs for suggestion in suggested)
+
+    # Those new suggestions also shouldn't be in the set of previously suggested configs.
+    assert all(suggestion.get_param_values() not in suggested_shuffled for suggestion in suggested)
+
+    grid_search_opt.bulk_register([suggestion.get_param_values() for suggestion in suggested],
+                                  [score] * len(suggested),
+                                  [status] * len(suggested))
+
+    assert all(suggestion.get_param_values() not in grid_search_opt.pending_configs for suggestion in suggested)
+    assert all(suggestion.get_param_values() not in grid_search_opt.suggested_configs for suggestion in suggested)
+
+    best_score, best_config = grid_search_opt.get_best_observation()
+    assert best_score == best_suggestion_score
+    assert best_config == best_suggestion