Make distribute_replications a required input argument (facebook#2051)

Summary: Pull Request resolved: facebook#2051 See title Reviewed By: saitcakmak, Balandat Differential Revision: D51431736 fbshipit-source-id: d6e2ee402d6b671d774afe89ff4c6e9b3e305d44
mgarrard · Dec 6, 2023 · 8481e87 · 8481e87
1 parent f76f721
commit 8481e87
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/ax/benchmark/methods/modular_botorch.py b/ax/benchmark/methods/modular_botorch.py
@@ -41,8 +41,8 @@
 def get_sobol_botorch_modular_acquisition(
     model_cls: Type[Model],
     acquisition_cls: Type[AcquisitionFunction],
+    distribute_replications: bool,
     scheduler_options: Optional[SchedulerOptions] = None,
-    distribute_replications: bool = False,
     name: Optional[str] = None,
 ) -> BenchmarkMethod:
     model_kwargs: Dict[

diff --git a/ax/benchmark/tests/test_benchmark.py b/ax/benchmark/tests/test_benchmark.py
@@ -147,6 +147,7 @@ def test_replication_mbm(self) -> None:
                     model_cls=SingleTaskGP,
                     acquisition_cls=qLogNoisyExpectedImprovement,
                     scheduler_options=get_sequential_optimization_scheduler_options(),
+                    distribute_replications=False,
                 ),
                 get_single_objective_benchmark_problem(infer_noise=False, num_trials=6),
                 "MBM::SingleTaskGP_qLogNEI",
@@ -155,6 +156,7 @@ def test_replication_mbm(self) -> None:
                 get_sobol_botorch_modular_acquisition(
                     model_cls=FixedNoiseGP,
                     acquisition_cls=qLogNoisyExpectedImprovement,
+                    distribute_replications=False,
                 ),
                 get_single_objective_benchmark_problem(infer_noise=False, num_trials=6),
                 "MBM::FixedNoiseGP_qLogNEI",
@@ -163,6 +165,7 @@ def test_replication_mbm(self) -> None:
                 get_sobol_botorch_modular_acquisition(
                     model_cls=FixedNoiseGP,
                     acquisition_cls=qNoisyExpectedHypervolumeImprovement,
+                    distribute_replications=False,
                 ),
                 get_multi_objective_benchmark_problem(infer_noise=False, num_trials=6),
                 "MBM::FixedNoiseGP_qNEHVI",
@@ -171,6 +174,7 @@ def test_replication_mbm(self) -> None:
                 get_sobol_botorch_modular_acquisition(
                     model_cls=SaasFullyBayesianSingleTaskGP,
                     acquisition_cls=qLogNoisyExpectedImprovement,
+                    distribute_replications=False,
                 ),
                 get_multi_objective_benchmark_problem(num_trials=6),
                 "MBM::SAAS_qLogNEI",
@@ -231,7 +235,9 @@ def test_benchmark_multiple_problems_methods(self) -> None:
             methods=[
                 get_sobol_benchmark_method(),
                 get_sobol_botorch_modular_acquisition(
-                    model_cls=SingleTaskGP, acquisition_cls=qLogNoisyExpectedImprovement
+                    model_cls=SingleTaskGP,
+                    acquisition_cls=qLogNoisyExpectedImprovement,
+                    distribute_replications=False,
                 ),
             ],
             seeds=(0, 1),
@@ -250,7 +256,9 @@ def test_timeout(self) -> None:
         )
 
         generation_strategy = get_sobol_botorch_modular_acquisition(
-            model_cls=SingleTaskGP, acquisition_cls=qLogNoisyExpectedImprovement
+            model_cls=SingleTaskGP,
+            acquisition_cls=qLogNoisyExpectedImprovement,
+            distribute_replications=False,
         ).generation_strategy
 
         method = BenchmarkMethod(

diff --git a/ax/benchmark/tests/test_methods.py b/ax/benchmark/tests/test_methods.py
@@ -24,6 +24,7 @@ def test_mbm_acquisition(self) -> None:
             model_cls=SingleTaskGP,
             acquisition_cls=qKnowledgeGradient,
             scheduler_options=get_sequential_optimization_scheduler_options(),
+            distribute_replications=False,
         )
         self.assertEqual(method.name, "MBM::SingleTaskGP_qKnowledgeGradient")
         gs = method.generation_strategy
@@ -44,6 +45,7 @@ def test_benchmark_replication_runs(self) -> None:
             model_cls=SingleTaskGP,
             scheduler_options=get_sequential_optimization_scheduler_options(),
             acquisition_cls=LogExpectedImprovement,
+            distribute_replications=False,
         )
         n_sobol_trials = method.generation_strategy._steps[0].num_trials
         # Only run one non-Sobol trial