Allow customizing num_init_trials in MBM benchmark method (#2286)

Summary:
Pull Request resolved: #2286

This is useful for benchmarking high parallelism setting, where we typically want all of the first batch to be Sobol.

Also makes use of the previously ignored name kwarg.

Reviewed By: esantorella

Differential Revision: D54690298

fbshipit-source-id: b65fc28df516402dd0114ee5dd6138998fac0d1c

ax/benchmark/methods/modular_botorch.py

@@ -14,7 +14,6 @@
 from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
 from ax.modelbridge.registry import Models
 from ax.models.torch.botorch_modular.model import SurrogateSpec
-
 from ax.service.scheduler import SchedulerOptions
 from botorch.acquisition.acquisition import AcquisitionFunction
 from botorch.acquisition.analytic import LogExpectedImprovement
@@ -59,13 +58,15 @@ def get_sobol_botorch_modular_acquisition(
     acqf_name = acqf_name_abbreviations.get(
         acquisition_cls.__name__, acquisition_cls.__name__
     )
-    name = f"MBM::{model_name}_{acqf_name}"
+    name = name or f"MBM::{model_name}_{acqf_name}"
 
     generation_strategy = GenerationStrategy(
         name=name,
         steps=[
             GenerationStep(
-                model=Models.SOBOL, num_trials=num_sobol_trials, min_trials_observed=5
+                model=Models.SOBOL,
+                num_trials=num_sobol_trials,
+                min_trials_observed=num_sobol_trials,
             ),
             GenerationStep(
                 model=Models.BOTORCH_MODULAR,

ax/benchmark/tests/methods/test_methods.py

@@ -48,9 +48,13 @@ def test_benchmark_replication_runs(self) -> None:
             model_cls=SingleTaskGP,
             scheduler_options=get_sequential_optimization_scheduler_options(),
             acquisition_cls=LogExpectedImprovement,
+            num_sobol_trials=2,
+            name="test",
             distribute_replications=False,
         )
         n_sobol_trials = method.generation_strategy._steps[0].num_trials
+        self.assertEqual(n_sobol_trials, 2)
+        self.assertEqual(method.name, "test")
         # Only run one non-Sobol trial
         problem = get_problem(problem_name="ackley4", num_trials=n_sobol_trials + 1)
         result = benchmark_replication(problem=problem, method=method, seed=0)