diff --git a/ax/benchmark/benchmark_method.py b/ax/benchmark/benchmark_method.py index 5e9cf8ff015..d12b0df675f 100644 --- a/ax/benchmark/benchmark_method.py +++ b/ax/benchmark/benchmark_method.py @@ -105,18 +105,6 @@ def get_benchmark_scheduler_options( ) -def get_sequential_optimization_scheduler_options( - timeout_hours: int = 4, -) -> SchedulerOptions: - """The typical SchedulerOptions used in benchmarking. - - Args: - timeout_hours: The maximum amount of time (in hours) to run each - benchmark replication. Defaults to 4 hours. - """ - return get_benchmark_scheduler_options(timeout_hours=timeout_hours) - - def _assign_frozen_attr(obj: Any, name: str, value: Any) -> None: # pyre-ignore [2] """Assign a new value to an attribute of a frozen dataclass. This is an ugly hack and shouldn't be used broadly. diff --git a/ax/benchmark/methods/modular_botorch.py b/ax/benchmark/methods/modular_botorch.py index 795065c9e17..379f92327aa 100644 --- a/ax/benchmark/methods/modular_botorch.py +++ b/ax/benchmark/methods/modular_botorch.py @@ -9,7 +9,7 @@ from ax.benchmark.benchmark_method import ( BenchmarkMethod, - get_sequential_optimization_scheduler_options, + get_benchmark_scheduler_options, ) from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy from ax.modelbridge.registry import Models @@ -89,7 +89,6 @@ def get_sobol_botorch_modular_acquisition( return BenchmarkMethod( name=generation_strategy.name, generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), + scheduler_options=scheduler_options or get_benchmark_scheduler_options(), distribute_replications=distribute_replications, ) diff --git a/ax/benchmark/methods/sobol.py b/ax/benchmark/methods/sobol.py index df828eb0d8a..6185ab1cc7c 100644 --- a/ax/benchmark/methods/sobol.py +++ b/ax/benchmark/methods/sobol.py @@ -9,7 +9,7 @@ from ax.benchmark.benchmark_method import ( BenchmarkMethod, - get_sequential_optimization_scheduler_options, + get_benchmark_scheduler_options, ) from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy from ax.modelbridge.registry import Models @@ -28,7 +28,6 @@ def get_sobol_benchmark_method( return BenchmarkMethod( name=generation_strategy.name, generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), + scheduler_options=scheduler_options or get_benchmark_scheduler_options(), distribute_replications=distribute_replications, ) diff --git a/ax/benchmark/tests/test_benchmark.py b/ax/benchmark/tests/test_benchmark.py index ab847249b0d..7dc66586cd0 100644 --- a/ax/benchmark/tests/test_benchmark.py +++ b/ax/benchmark/tests/test_benchmark.py @@ -18,7 +18,7 @@ ) from ax.benchmark.benchmark_method import ( BenchmarkMethod, - get_sequential_optimization_scheduler_options, + get_benchmark_scheduler_options, ) from ax.benchmark.benchmark_problem import SingleObjectiveBenchmarkProblem from ax.benchmark.benchmark_result import BenchmarkResult @@ -302,7 +302,7 @@ def test_replication_mbm(self) -> None: get_sobol_botorch_modular_acquisition( model_cls=SingleTaskGP, acquisition_cls=qLogNoisyExpectedImprovement, - scheduler_options=get_sequential_optimization_scheduler_options(), + scheduler_options=get_benchmark_scheduler_options(), distribute_replications=False, ), get_single_objective_benchmark_problem( diff --git a/ax/benchmark/tests/test_benchmark_method.py b/ax/benchmark/tests/test_benchmark_method.py index bd556b143e3..f6ab6c29f29 100644 --- a/ax/benchmark/tests/test_benchmark_method.py +++ b/ax/benchmark/tests/test_benchmark_method.py @@ -7,7 +7,7 @@ from ax.benchmark.benchmark_method import ( BenchmarkMethod, - get_sequential_optimization_scheduler_options, + get_benchmark_scheduler_options, ) from ax.modelbridge.generation_strategy import ( GenerationNode, @@ -26,7 +26,7 @@ def test_benchmark_method(self) -> None: steps=[GenerationStep(model=Models.SOBOL, num_trials=10)], name="SOBOL", ) - options = get_sequential_optimization_scheduler_options() + options = get_benchmark_scheduler_options() method = BenchmarkMethod( name="Sobol10", generation_strategy=gs, scheduler_options=options ) @@ -41,7 +41,7 @@ def test_benchmark_method(self) -> None: self.assertEqual(options.min_seconds_before_poll, 0) self.assertEqual(options.timeout_hours, 4) - options = get_sequential_optimization_scheduler_options(timeout_hours=10) + options = get_benchmark_scheduler_options(timeout_hours=10) method = BenchmarkMethod( name="Sobol10", generation_strategy=gs, scheduler_options=options )