Add e2e tests with HSS (facebook#2345)

Summary:
Pull Request resolved: facebook#2345

Adds tests for candidate generation, prediction & cross validation using HSS with MBM defaults. Notably, the tests do not involve any mocks, except for `fast_botorch_optimize`.

The tests also highlight a weakness of current setup as it does not allow making predictions using valid parameterizations, when those parameterizations lack the  inactive parameters (i.e. they're not full parmeterizations).

Reviewed By: Balandat

Differential Revision: D55952289

fbshipit-source-id: 90dbb923d387d05ef05af4e013e6cec38174f8e7

Author: saitcakmak

ax/modelbridge/tests/test_hierarchical_search_space.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from contextlib import ExitStack
+from random import random
+from typing import List
+
+from ax.core.experiment import Experiment
+from ax.core.objective import Objective
+from ax.core.observation import ObservationFeatures
+from ax.core.optimization_config import OptimizationConfig
+from ax.core.parameter import (
+    ChoiceParameter,
+    FixedParameter,
+    ParameterType,
+    RangeParameter,
+)
+from ax.core.search_space import HierarchicalSearchSpace
+from ax.core.trial import Trial
+from ax.metrics.noisy_function import GenericNoisyFunctionMetric
+from ax.modelbridge.cross_validation import cross_validate
+from ax.modelbridge.registry import Models
+from ax.runners.synthetic import SyntheticRunner
+from ax.utils.common.constants import Keys
+from ax.utils.common.testutils import TestCase
+from ax.utils.common.typeutils import checked_cast, not_none
+from ax.utils.testing.mock import fast_botorch_optimize
+
+
+class TestHierarchicalSearchSpace(TestCase):
+    """Tests for various modelbridge functionality with commonly used transforms
+    using hierarchical search spaces (HSS).
+    """
+
+    def setUp(self) -> None:
+        int_range = RangeParameter(
+            name="int_range",
+            parameter_type=ParameterType.INT,
+            lower=0,
+            upper=10,
+        )
+        str_choice = ChoiceParameter(
+            name="str_choice",
+            parameter_type=ParameterType.STRING,
+            values=["a", "b", "c"],
+        )
+        fixed_root = FixedParameter(
+            name="root",
+            parameter_type=ParameterType.STRING,
+            value="root",
+            dependents={"root": ["int_range", "str_choice"]},
+        )
+        # This HSS does not have a real hierarchy.
+        self.non_hierarchical_hss = HierarchicalSearchSpace(
+            parameters=[
+                fixed_root,
+                int_range,
+                str_choice,
+            ]
+        )
+        choice_root = ChoiceParameter(
+            name="root",
+            parameter_type=ParameterType.STRING,
+            values=["range", "choice"],
+            dependents={"range": ["int_range"], "choice": ["str_choice"]},
+        )
+        # This HSS has a simple hierarchy -- one parameter on each branch.
+        self.simple_hss = HierarchicalSearchSpace(
+            parameters=[choice_root, int_range, str_choice]
+        )
+        fixed_leaf = FixedParameter(
+            name="fixed_leaf",
+            parameter_type=ParameterType.STRING,
+            value="leaf",
+        )
+        middle_choice = ChoiceParameter(
+            name="middle_choice",
+            parameter_type=ParameterType.INT,
+            values=[0, 1],
+            dependents={0: ["fixed_leaf"], 1: ["int_range", "str_choice"]},
+        )
+        int_choice = ChoiceParameter(
+            name="int_choice",
+            parameter_type=ParameterType.INT,
+            values=[0, 1, 2, 3],
+            is_ordered=False,
+        )
+        float_range = RangeParameter(
+            name="float_range",
+            parameter_type=ParameterType.FLOAT,
+            lower=0.0,
+            upper=5.0,
+        )
+        choice_root2 = ChoiceParameter(
+            name="root2",
+            parameter_type=ParameterType.BOOL,
+            values=[True, False],
+            dependents={True: ["middle_choice", "float_range"], False: ["int_choice"]},
+        )
+        # This HSS has a more complex, multi-level hierarchy.
+        self.complex_hss = HierarchicalSearchSpace(
+            parameters=[
+                choice_root2,
+                int_choice,
+                middle_choice,
+                float_range,
+                fixed_leaf,
+                int_range,
+                str_choice,
+            ]
+        )
+
+    @fast_botorch_optimize
+    def _test_gen_base(
+        self,
+        hss: HierarchicalSearchSpace,
+        expected_num_candidate_params: List[int],
+        num_sobol_trials: int = 5,
+        num_bo_trials: int = 5,
+    ) -> Experiment:
+        """Test Sobol & MBM candidate generation with HSS using default transforms.
+
+        Args:
+            hss: The hierarchical search space to test.
+            expected_num_candidate_params: The expected number of parameters in each
+                candidate. This list should include all possible values, since different
+                branches of HSS may have different numbers of parameters.
+            num_sobol_trials: The number of Sobol trials to run.
+            num_bo_trials: The number of BO trials to run.
+
+        Returns:
+            The experiment with the generated candidates. This can be used to chain
+            tests for other functionality that requires data.
+        """
+        experiment = Experiment(
+            name="test_experiment",
+            search_space=hss,
+            optimization_config=OptimizationConfig(
+                objective=Objective(
+                    metric=GenericNoisyFunctionMetric(
+                        name="random", f=lambda _: random()
+                    ),
+                    minimize=True,
+                )
+            ),
+            runner=SyntheticRunner(),
+        )
+
+        sobol = Models.SOBOL(search_space=hss)
+        for _ in range(num_sobol_trials):
+            trial = experiment.new_trial(generator_run=sobol.gen(n=1))
+            trial.run().mark_completed()
+
+        for _ in range(num_bo_trials):
+            mbm = Models.BOTORCH_MODULAR(
+                experiment=experiment, data=experiment.fetch_data()
+            )
+            trial = experiment.new_trial(generator_run=mbm.gen(n=1))
+            trial.run().mark_completed()
+
+        for t in experiment.trials.values():
+            trial = checked_cast(Trial, t)
+            arm = not_none(trial.arm)
+            self.assertIn(len(arm.parameters), expected_num_candidate_params)
+            # Check that the trials have the full parameterization recorded.
+            full_parameterization = not_none(
+                trial._get_candidate_metadata(arm_name=arm.name)
+            )[Keys.FULL_PARAMETERIZATION]
+            self.assertEqual(full_parameterization.keys(), hss.parameters.keys())
+
+        return experiment
+
+    @fast_botorch_optimize
+    def _base_test_predict_and_cv(
+        self,
+        experiment: Experiment,
+        expect_errors_with_final_parameterization: bool = False,
+    ) -> None:
+        """Test predict and cross validation with a given experiment.
+        The predict tests are done using the full parameterization, the
+        final parameterization with the full parameterization recorded in
+        metadata, and with the final parameterization only. When the final
+        parameterization lacks some parameters, this may error out.
+        `expect_errors_with_final_parameterization` arg is used to handle
+        the `KeyError` that is expected (but should be fixed) in this setting.
+        """
+        mbm = Models.BOTORCH_MODULAR(
+            experiment=experiment, data=experiment.fetch_data()
+        )
+        for t in experiment.trials.values():
+            trial = checked_cast(Trial, t)
+            arm = not_none(trial.arm)
+            final_parameterization = arm.parameters
+            full_parameterization = not_none(
+                trial._get_candidate_metadata(arm_name=arm.name)
+            )[Keys.FULL_PARAMETERIZATION]
+            # Predict with full parameterization -- this should always work.
+            mbm.predict([ObservationFeatures(parameters=full_parameterization)])
+            # Predict with final parameterization -- this may error out :(.
+            with ExitStack() as es:
+                if expect_errors_with_final_parameterization:
+                    es.enter_context(self.assertRaises(KeyError))
+                mbm.predict([ObservationFeatures(parameters=final_parameterization)])
+            # Predict with final parameterization but include the full parameterization
+            # in the metadata. This is similar to what happens inside cross_validate.
+            mbm.predict(
+                [
+                    ObservationFeatures(
+                        parameters=final_parameterization,
+                        metadata={Keys.FULL_PARAMETERIZATION: full_parameterization},
+                    )
+                ]
+            )
+        cv_res = cross_validate(model=mbm)
+        self.assertEqual(len(cv_res), len(experiment.trials))
+
+    def test_with_non_hierarchical_hss(self) -> None:
+        experiment = self._test_gen_base(
+            hss=self.non_hierarchical_hss, expected_num_candidate_params=[3]
+        )
+        self._base_test_predict_and_cv(experiment=experiment)
+
+    def test_with_simple_hss(self) -> None:
+        experiment = self._test_gen_base(
+            hss=self.simple_hss, expected_num_candidate_params=[2]
+        )
+        self._base_test_predict_and_cv(
+            experiment=experiment, expect_errors_with_final_parameterization=True
+        )
+
+    def test_with_complex_hss(self) -> None:
+        experiment = self._test_gen_base(
+            hss=self.complex_hss, expected_num_candidate_params=[2, 4, 5]
+        )
+        self._base_test_predict_and_cv(
+            experiment=experiment, expect_errors_with_final_parameterization=True
+        )