Replace Models.GPEI usage with Models.BoTorchModular (#2733)

Summary:
Pull Request resolved: #2733

`Models.GPEI` points to the deprecated legacy Ax model. For most practical usage, it behaves identical to `Models.BOTORCH_MODULAR` (aka MBM), which is the recommmended model to use in Ax. This diff updates many use cases to make it easier to fully deprecate the registry entry down the line.

Reviewed By: Balandat

Differential Revision: D62146528

fbshipit-source-id: 7f53f9748432b70c477690d20959e21809122435

ax/modelbridge/registry.py

@@ -424,7 +424,7 @@ class Models(ModelRegistryBase):
 
     To instantiate a model in this enum, simply call an enum member like so:
     `Models.SOBOL(search_space=search_space)` or
-    `Models.GPEI(experiment=experiment, data=data)`. Keyword arguments
+    `Models.BOTORCH(experiment=experiment, data=data)`. Keyword arguments
     specified to the call will be passed into the model or the model bridge
     constructors according to their keyword.
 

ax/modelbridge/tests/test_aepsych_criterion.py

@@ -31,13 +31,13 @@ def test_single_criterion(self) -> None:
         experiment = get_experiment()
 
         generation_strategy = GenerationStrategy(
-            name="SOBOL+GPEI::default",
+            name="SOBOL+MBM::default",
             steps=[
                 GenerationStep(
                     model=Models.SOBOL, num_trials=-1, completion_criteria=[criterion]
                 ),
                 GenerationStep(
-                    model=Models.GPEI,
+                    model=Models.BOTORCH_MODULAR,
                     num_trials=-1,
                     max_parallelism=1,
                 ),
@@ -77,7 +77,8 @@ def test_single_criterion(self) -> None:
             )
 
             self.assertEqual(
-                generation_strategy._curr.model_spec_to_gen_from.model_enum, Models.GPEI
+                generation_strategy._curr.model_spec_to_gen_from.model_enum,
+                Models.BOTORCH_MODULAR,
             )
 
     def test_many_criteria(self) -> None:
@@ -89,13 +90,13 @@ def test_many_criteria(self) -> None:
         experiment = get_experiment()
 
         generation_strategy = GenerationStrategy(
-            name="SOBOL+GPEI::default",
+            name="SOBOL+MBM::default",
             steps=[
                 GenerationStep(
                     model=Models.SOBOL, num_trials=-1, completion_criteria=criteria
                 ),
                 GenerationStep(
-                    model=Models.GPEI,
+                    model=Models.BOTORCH_MODULAR,
                     num_trials=-1,
                     max_parallelism=1,
                 ),
@@ -153,5 +154,6 @@ def test_many_criteria(self) -> None:
             )
 
             self.assertEqual(
-                generation_strategy._curr.model_spec_to_gen_from.model_enum, Models.GPEI
+                generation_strategy._curr.model_spec_to_gen_from.model_enum,
+                Models.BOTORCH_MODULAR,
             )

ax/modelbridge/tests/test_generation_node.py

@@ -116,7 +116,7 @@ def test_properties(self) -> None:
             node_name="test",
             model_specs=[
                 ModelSpec(
-                    model_enum=Models.GPEI,
+                    model_enum=Models.BOTORCH_MODULAR,
                     model_kwargs={},
                     model_gen_kwargs={
                         "n": 1,
@@ -140,7 +140,7 @@ def test_properties(self) -> None:
         self.assertEqual(
             node.model_spec_to_gen_from.model_kwargs, node.model_specs[0].model_kwargs
         )
-        self.assertEqual(node.model_to_gen_from_name, "GPEI")
+        self.assertEqual(node.model_to_gen_from_name, "BoTorch")
         self.assertEqual(
             node.model_spec_to_gen_from.model_gen_kwargs,
             node.model_specs[0].model_gen_kwargs,
@@ -167,7 +167,7 @@ def test_node_string_representation(self) -> None:
             node_name="test",
             model_specs=[
                 ModelSpec(
-                    model_enum=Models.GPEI,
+                    model_enum=Models.BOTORCH_MODULAR,
                     model_kwargs={},
                     model_gen_kwargs={},
                 ),
@@ -180,7 +180,7 @@ def test_node_string_representation(self) -> None:
         self.assertEqual(
             string_rep,
             (
-                "GenerationNode(model_specs=[ModelSpec(model_enum=GPEI,"
+                "GenerationNode(model_specs=[ModelSpec(model_enum=BoTorch,"
                 " model_kwargs={}, model_gen_kwargs={}, model_cv_kwargs={},"
                 " )], node_name=test, "
                 "transition_criteria=[MaxTrials({'threshold': 5, "
@@ -197,7 +197,7 @@ def test_single_fixed_features(self) -> None:
             node_name="test",
             model_specs=[
                 ModelSpec(
-                    model_enum=Models.GPEI,
+                    model_enum=Models.BOTORCH_MODULAR,
                     model_kwargs={},
                     model_gen_kwargs={
                         "n": 2,

ax/modelbridge/tests/test_model_fit_metrics.py

@@ -58,7 +58,7 @@ def setUp(self) -> None:
                 GenerationStep(
                     model=Models.SOBOL, num_trials=NUM_SOBOL, max_parallelism=NUM_SOBOL
                 ),
-                GenerationStep(model=Models.GPEI, num_trials=-1),
+                GenerationStep(model=Models.BOTORCH_MODULAR, num_trials=-1),
             ]
         )
 

ax/modelbridge/tests/test_model_spec.py

@@ -37,7 +37,7 @@ def setUp(self) -> None:
 class ModelSpecTest(BaseModelSpecTest):
     @fast_botorch_optimize
     def test_construct(self) -> None:
-        ms = ModelSpec(model_enum=Models.GPEI)
+        ms = ModelSpec(model_enum=Models.BOTORCH_MODULAR)
         with self.assertRaises(UserInputError):
             ms.gen(n=1)
         ms.fit(experiment=self.experiment, data=self.data)
@@ -51,7 +51,7 @@ def test_construct(self) -> None:
         wraps=extract_search_space_digest,
     )
     def test_fit(self, wrapped_extract_ssd: Mock) -> None:
-        ms = ModelSpec(model_enum=Models.GPEI)
+        ms = ModelSpec(model_enum=Models.BOTORCH_MODULAR)
         # This should fit the model as usual.
         ms.fit(experiment=self.experiment, data=self.data)
         wrapped_extract_ssd.assert_called_once()
@@ -159,7 +159,7 @@ def test_cross_validate_with_non_GP_model(
         mock_diagnostics.assert_not_called()
 
     def test_fixed_features(self) -> None:
-        ms = ModelSpec(model_enum=Models.GPEI)
+        ms = ModelSpec(model_enum=Models.BOTORCH_MODULAR)
         self.assertIsNone(ms.fixed_features)
         new_features = ObservationFeatures(parameters={"a": 1.0})
         ms.fixed_features = new_features
@@ -177,7 +177,7 @@ def test_gen_attaches_empty_model_fit_metadata_if_fit_not_applicable(self) -> No
         self.assertEqual(gen_metadata["model_std_generalization"], None)
 
     def test_gen_attaches_model_fit_metadata_if_applicable(self) -> None:
-        ms = ModelSpec(model_enum=Models.GPEI)
+        ms = ModelSpec(model_enum=Models.BOTORCH_MODULAR)
         ms.fit(experiment=self.experiment, data=self.data)
         gr = ms.gen(n=1)
         gen_metadata = not_none(gr.gen_metadata)

ax/modelbridge/tests/test_transition_criterion.py

@@ -58,7 +58,7 @@ def test_minimum_preference_criterion(self) -> None:
                     completion_criteria=[criterion],
                 ),
                 GenerationStep(
-                    model=Models.GPEI,
+                    model=Models.BOTORCH_MODULAR,
                     num_trials=-1,
                     max_parallelism=1,
                 ),
@@ -92,7 +92,8 @@ def test_minimum_preference_criterion(self) -> None:
                 )
             )
             self.assertEqual(
-                generation_strategy._curr.model_spec_to_gen_from.model_enum, Models.GPEI
+                generation_strategy._curr.model_spec_to_gen_from.model_enum,
+                Models.BOTORCH_MODULAR,
             )
 
     def test_default_step_criterion_setup(self) -> None:
@@ -106,21 +107,21 @@ def test_default_step_criterion_setup(self) -> None:
         """
         experiment = get_experiment()
         gs = GenerationStrategy(
-            name="SOBOL+GPEI::default",
+            name="SOBOL+MBM::default",
             steps=[
                 GenerationStep(
                     model=Models.SOBOL,
                     num_trials=3,
                 ),
                 GenerationStep(
-                    model=Models.GPEI,
+                    model=Models.BOTORCH_MODULAR,
                     num_trials=4,
                     max_parallelism=1,
                     min_trials_observed=2,
                     enforce_num_trials=False,
                 ),
                 GenerationStep(
-                    model=Models.GPEI,
+                    model=Models.BOTORCH_MODULAR,
                     num_trials=-1,
                 ),
             ],

ax/plot/tests/test_traces.py

@@ -77,7 +77,7 @@ def test_plot_objective_value_vs_trial_index(self) -> None:
         for _ in range(2):
             t = exp.new_trial(sobol.gen(1)).run()
             t.mark_completed()
-        model = Models.GPEI(
+        model = Models.BOTORCH_MODULAR(
             experiment=exp,
             data=exp.fetch_data(),
         )

ax/service/tests/test_report_utils.py

@@ -522,7 +522,7 @@ def test_skip_contour_high_dimensional(self) -> None:
         sobol = Models.SOBOL(search_space=exp.search_space)
         for _ in range(1):
             exp.new_trial(sobol.gen(1)).run()
-        model = Models.GPEI(
+        model = Models.BOTORCH_MODULAR(
             experiment=exp,
             data=exp.fetch_data(),
         )
@@ -1206,7 +1206,7 @@ def test_compare_to_baseline_select_baseline_arm(self) -> None:
                 experiment=experiment,
                 arms_df=arms_df,
                 baseline_arm_name=wrong_baseline_name,
-            ),
+            )
 
         # status quo baseline arm
         experiment_with_status_quo = copy.deepcopy(experiment)
@@ -1272,7 +1272,9 @@ def test_warn_if_unpredictable_metrics(self) -> None:
                     min_trials_observed=3,
                     max_parallelism=3,
                 ),
-                GenerationStep(model=Models.GPEI, num_trials=-1, max_parallelism=3),
+                GenerationStep(
+                    model=Models.BOTORCH_MODULAR, num_trials=-1, max_parallelism=3
+                ),
             ]
         )
         gs.experiment = exp

ax/telemetry/tests/test_scheduler.py

@@ -167,7 +167,7 @@ def test_scheduler_model_fit_metrics_logging(self) -> None:
                 GenerationStep(
                     model=Models.SOBOL, num_trials=NUM_SOBOL, max_parallelism=NUM_SOBOL
                 ),
-                GenerationStep(model=Models.GPEI, num_trials=-1),
+                GenerationStep(model=Models.BOTORCH_MODULAR, num_trials=-1),
             ]
         )
 

docs/api.md

@@ -155,7 +155,7 @@ for i in range(5):
 
 best_arm = None
 for i in range(15):
-    gpei = Models.GPEI(experiment=exp, data=exp.fetch_data())
+    gpei = Models.BOTORCH_MODULAR(experiment=exp, data=exp.fetch_data())
     generator_run = gpei.gen(1)
     best_arm, _ = generator_run.best_arm_predictions
     trial = exp.new_trial(generator_run=generator_run)

docs/trial-evaluation.md

@@ -122,7 +122,7 @@ for i in range(5):
     trial.mark_completed()
 
 for i in range(15):
-    gpei = Models.GPEI(experiment=exp, data=exp.fetch_data())
+    gpei = Models.BOTORCH_MODULAR(experiment=exp, data=exp.fetch_data())
     generator_run = gpei.gen(1)
     trial = exp.new_trial(generator_run=generator_run)
     trial.run()