Get rid of SingleObjectiveBenchmarkProblem (#2606)

Summary: Pull Request resolved: #2606 The class is not adding anything. We should check if problems are single-objective by looking at the type of their `optimization_config`. This PR: * Removes `SingleObjectiveBenchmarkProblem` * Replaces all references to it with `BenchmarkProblem` * Removes `get_benchmark_problem`, which is redundant with `get_single_objective_benchmark_problem` Reviewed By: dme65 Differential Revision: D60285191
facebook · Jul 28, 2024 · 9d19074 · 9d19074
1 parent 79627e7
commit 9d19074
Show file tree

Hide file tree

Showing 9 changed files with 25 additions and 74 deletions.
diff --git a/ax/benchmark/benchmark_problem.py b/ax/benchmark/benchmark_problem.py
@@ -93,23 +93,18 @@ class BenchmarkProblem(Base):
     is_noiseless: bool
 
 
-class SingleObjectiveBenchmarkProblem(BenchmarkProblem):
-    """A `BenchmarkProblem` that supports a single objective."""
-
-    pass
-
-
 def create_single_objective_problem_from_botorch(
     test_problem_class: Type[SyntheticTestFunction],
     test_problem_kwargs: Dict[str, Any],
     lower_is_better: bool,
     num_trials: int,
     observe_noise_sd: bool = False,
-) -> SingleObjectiveBenchmarkProblem:
+) -> BenchmarkProblem:
     """
-    Create a BenchmarkProblem from a BoTorch BaseTestProblem using
-    specialized Metrics and Runners. The test problem's result will be
-    computed on the Runner and retrieved by the Metric.
+    Create a `BenchmarkProblem` whose `optimization_config` is a
+    `SingleObjectiveOptimizationConfig` a BoTorch SyntheticTestFunction using
+    specialized Metrics and Runners for benchmarking. The test problem's result
+    will be computed on the Runner and retrieved by the Metric.
 
     Args:
         test_problem_class: The BoTorch test problem class which will be used
@@ -188,7 +183,7 @@ def create_single_objective_problem_from_botorch(
         if isinstance(test_problem, MultiObjectiveTestProblem)
         else test_problem.optimal_value
     )
-    return SingleObjectiveBenchmarkProblem(
+    return BenchmarkProblem(
         name=name,
         search_space=search_space,
         optimization_config=optimization_config,

diff --git a/ax/benchmark/problems/hpo/pytorch_cnn.py b/ax/benchmark/problems/hpo/pytorch_cnn.py
@@ -9,7 +9,7 @@
 
 import pandas as pd
 import torch
-from ax.benchmark.benchmark_problem import SingleObjectiveBenchmarkProblem
+from ax.benchmark.benchmark_problem import BenchmarkProblem
 from ax.core.base_trial import BaseTrial, TrialStatus
 from ax.core.data import Data
 from ax.core.metric import Metric, MetricFetchE, MetricFetchResult
@@ -26,7 +26,7 @@
 from torch.utils.data import DataLoader, Dataset
 
 
-class PyTorchCNNBenchmarkProblem(SingleObjectiveBenchmarkProblem):
+class PyTorchCNNBenchmarkProblem(BenchmarkProblem):
     @equality_typechecker
     def __eq__(self, other: Base) -> bool:
         if not isinstance(other, PyTorchCNNBenchmarkProblem):

diff --git a/ax/benchmark/problems/synthetic/discretized/mixed_integer.py b/ax/benchmark/problems/synthetic/discretized/mixed_integer.py
@@ -20,7 +20,7 @@
 
 from typing import Dict, List, Optional, Tuple, Type, Union
 
-from ax.benchmark.benchmark_problem import SingleObjectiveBenchmarkProblem
+from ax.benchmark.benchmark_problem import BenchmarkProblem
 from ax.benchmark.metrics.benchmark import BenchmarkMetric
 from ax.benchmark.runners.botorch_test import BotorchTestProblemRunner
 from ax.core.objective import Objective
@@ -47,7 +47,7 @@ def _get_problem_from_common_inputs(
     num_trials: int,
     optimal_value: float,
     test_problem_bounds: Optional[List[Tuple[float, float]]] = None,
-) -> SingleObjectiveBenchmarkProblem:
+) -> BenchmarkProblem:
     """This is a helper that deduplicates common bits of the below problems.
 
     Args:
@@ -111,7 +111,7 @@ def _get_problem_from_common_inputs(
         outcome_names=[metric_name],
         modified_bounds=bounds,
     )
-    return SingleObjectiveBenchmarkProblem(
+    return BenchmarkProblem(
         name=benchmark_name + ("_observed_noise" if observe_noise_sd else ""),
         search_space=search_space,
         optimization_config=optimization_config,
@@ -128,7 +128,7 @@ def get_discrete_hartmann(
     num_trials: int = 50,
     observe_noise_sd: bool = False,
     bounds: Optional[List[Tuple[float, float]]] = None,
-) -> SingleObjectiveBenchmarkProblem:
+) -> BenchmarkProblem:
     """6D Hartmann problem where first 4 dimensions are discretized."""
     dim_int = 4
     if bounds is None:
@@ -160,7 +160,7 @@ def get_discrete_ackley(
     num_trials: int = 50,
     observe_noise_sd: bool = False,
     bounds: Optional[List[Tuple[float, float]]] = None,
-) -> SingleObjectiveBenchmarkProblem:
+) -> BenchmarkProblem:
     """13D Ackley problem where first 10 dimensions are discretized.
 
     This also restricts Ackley evaluation bounds to [0, 1].
@@ -193,7 +193,7 @@ def get_discrete_rosenbrock(
     num_trials: int = 50,
     observe_noise_sd: bool = False,
     bounds: Optional[List[Tuple[float, float]]] = None,
-) -> SingleObjectiveBenchmarkProblem:
+) -> BenchmarkProblem:
     """10D Rosenbrock problem where first 6 dimensions are discretized."""
     dim_int = 6
     if bounds is None:

diff --git a/ax/benchmark/problems/synthetic/hss/jenatton.py b/ax/benchmark/problems/synthetic/hss/jenatton.py
@@ -5,7 +5,7 @@
 
 # pyre-strict
 
-from ax.benchmark.benchmark_problem import SingleObjectiveBenchmarkProblem
+from ax.benchmark.benchmark_problem import BenchmarkProblem
 from ax.benchmark.metrics.jenatton import JenattonMetric
 from ax.core.objective import Objective
 from ax.core.optimization_config import OptimizationConfig
@@ -17,7 +17,7 @@
 def get_jenatton_benchmark_problem(
     num_trials: int = 50,
     observe_noise_sd: bool = False,
-) -> SingleObjectiveBenchmarkProblem:
+) -> BenchmarkProblem:
     search_space = HierarchicalSearchSpace(
         parameters=[
             ChoiceParameter(
@@ -65,7 +65,7 @@ def get_jenatton_benchmark_problem(
 
     name = "Jenatton" + ("_observed_noise" if observe_noise_sd else "")
 
-    return SingleObjectiveBenchmarkProblem(
+    return BenchmarkProblem(
         name=name,
         search_space=search_space,
         optimization_config=optimization_config,

diff --git a/ax/benchmark/tests/test_benchmark_problem.py b/ax/benchmark/tests/test_benchmark_problem.py
@@ -85,7 +85,7 @@ def test_single_objective_from_botorch(self) -> None:
                     test_problem.optimization_config.outcome_constraints, []
                 )
                 expected_repr = (
-                    "SingleObjectiveBenchmarkProblem(name='Ackley', "
+                    "BenchmarkProblem(name='Ackley', "
                     "optimization_config=OptimizationConfig(objective=Objective("
                     'metric_name="Ackley", '
                     "minimize=True), outcome_constraints=[]), "
@@ -105,7 +105,7 @@ def test_single_objective_from_botorch(self) -> None:
                 self.assertFalse(outcome_constraint.relative)
                 self.assertEqual(outcome_constraint.bound, 0.0)
                 expected_repr = (
-                    "SingleObjectiveBenchmarkProblem(name='ConstrainedHartmann', "
+                    "BenchmarkProblem(name='ConstrainedHartmann', "
                     "optimization_config=OptimizationConfig(objective=Objective("
                     'metric_name="ConstrainedHartmann", minimize=True), '
                     "outcome_constraints=[OutcomeConstraint(constraint_slack_0"

diff --git a/ax/storage/json_store/encoders.py b/ax/storage/json_store/encoders.py
@@ -14,7 +14,6 @@
 from ax.benchmark.benchmark_problem import (
     BenchmarkProblem,
     MultiObjectiveBenchmarkProblem,
-    SingleObjectiveBenchmarkProblem,
 )
 from ax.benchmark.problems.hpo.torchvision import PyTorchCNNTorchvisionBenchmarkProblem
 from ax.core import ObservationFeatures
@@ -155,38 +154,9 @@ def multi_objective_benchmark_problem_to_dict(
     moo_benchmark_problem: MultiObjectiveBenchmarkProblem,
 ) -> Dict[str, Any]:
     """Convert Ax multi-objective benchmark problem to a dictionary."""
-    return {
-        "__type": moo_benchmark_problem.__class__.__name__,
-        "name": moo_benchmark_problem.name,
-        "search_space": moo_benchmark_problem.search_space,
-        "optimization_config": moo_benchmark_problem.optimization_config,
-        "runner": moo_benchmark_problem.runner,
-        "num_trials": moo_benchmark_problem.num_trials,
-        "is_noiseless": moo_benchmark_problem.is_noiseless,
-        "observe_noise_stds": moo_benchmark_problem.observe_noise_stds,
-        "has_ground_truth": moo_benchmark_problem.has_ground_truth,
-        "tracking_metrics": moo_benchmark_problem.tracking_metrics,
-        "optimal_value": moo_benchmark_problem.optimal_value,
-        "reference_point": moo_benchmark_problem.reference_point,
-    }
-
-
-def single_objective_benchmark_problem_to_dict(
-    soo_benchmark_problem: SingleObjectiveBenchmarkProblem,
-) -> Dict[str, Any]:
-    return {
-        "__type": soo_benchmark_problem.__class__.__name__,
-        "name": soo_benchmark_problem.name,
-        "search_space": soo_benchmark_problem.search_space,
-        "optimization_config": soo_benchmark_problem.optimization_config,
-        "runner": soo_benchmark_problem.runner,
-        "num_trials": soo_benchmark_problem.num_trials,
-        "is_noiseless": soo_benchmark_problem.is_noiseless,
-        "observe_noise_stds": soo_benchmark_problem.observe_noise_stds,
-        "has_ground_truth": soo_benchmark_problem.has_ground_truth,
-        "tracking_metrics": soo_benchmark_problem.tracking_metrics,
-        "optimal_value": soo_benchmark_problem.optimal_value,
-    }
+    result = benchmark_problem_to_dict(moo_benchmark_problem)
+    result["reference_point"] = moo_benchmark_problem.reference_point
+    return result
 
 
 def trial_to_dict(trial: Trial) -> Dict[str, Any]:

diff --git a/ax/storage/json_store/registry.py b/ax/storage/json_store/registry.py
@@ -14,7 +14,6 @@
 from ax.benchmark.benchmark_problem import (
     BenchmarkProblem,
     MultiObjectiveBenchmarkProblem,
-    SingleObjectiveBenchmarkProblem,
 )
 from ax.benchmark.benchmark_result import AggregatedBenchmarkResult, BenchmarkResult
 from ax.benchmark.metrics.benchmark import BenchmarkMetric, GroundTruthBenchmarkMetric
@@ -157,7 +156,6 @@
     runner_to_dict,
     scalarized_objective_to_dict,
     search_space_to_dict,
-    single_objective_benchmark_problem_to_dict,
     sum_parameter_constraint_to_dict,
     surrogate_to_dict,
     threshold_early_stopping_strategy_to_dict,
@@ -256,7 +254,6 @@
     ScalarizedObjective: scalarized_objective_to_dict,
     SearchSpace: search_space_to_dict,
     SingleDiagnosticBestModelSelector: best_model_selector_to_dict,
-    SingleObjectiveBenchmarkProblem: single_objective_benchmark_problem_to_dict,
     HierarchicalSearchSpace: search_space_to_dict,
     SumConstraint: sum_parameter_constraint_to_dict,
     Surrogate: surrogate_to_dict,
@@ -382,7 +379,6 @@
     "SchedulerOptions": SchedulerOptions,
     "SearchSpace": SearchSpace,
     "SingleDiagnosticBestModelSelector": SingleDiagnosticBestModelSelector,
-    "SingleObjectiveBenchmarkProblem": SingleObjectiveBenchmarkProblem,
     "SklearnDataset": SklearnDataset,
     "SklearnMetric": SklearnMetric,
     "SklearnModelType": SklearnModelType,

diff --git a/ax/storage/json_store/tests/test_json_store.py b/ax/storage/json_store/tests/test_json_store.py
@@ -46,7 +46,6 @@
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.benchmark_stubs import (
     get_aggregated_benchmark_result,
-    get_benchmark_problem,
     get_benchmark_result,
     get_multi_objective_benchmark_problem,
     get_single_objective_benchmark_problem,
@@ -142,7 +141,7 @@
     ("AugmentedHartmannMetric", get_augmented_hartmann_metric),
     ("BatchTrial", get_batch_trial),
     ("BenchmarkMethod", get_sobol_gpei_benchmark_method),
-    ("BenchmarkProblem", get_benchmark_problem),
+    ("BenchmarkProblem", get_single_objective_benchmark_problem),
     ("BenchmarkResult", get_benchmark_result),
     ("BoTorchModel", get_botorch_model),
     ("BoTorchModel", get_botorch_model_with_default_acquisition_class),

diff --git a/ax/utils/testing/benchmark_stubs.py b/ax/utils/testing/benchmark_stubs.py
@@ -11,10 +11,10 @@
 import numpy as np
 from ax.benchmark.benchmark_method import BenchmarkMethod
 from ax.benchmark.benchmark_problem import (
+    BenchmarkProblem,
     create_multi_objective_problem_from_botorch,
     create_single_objective_problem_from_botorch,
     MultiObjectiveBenchmarkProblem,
-    SingleObjectiveBenchmarkProblem,
 )
 from ax.benchmark.benchmark_result import AggregatedBenchmarkResult, BenchmarkResult
 from ax.benchmark.problems.surrogate import (
@@ -45,20 +45,11 @@
 from botorch.test_functions.synthetic import Branin
 
 
-def get_benchmark_problem() -> SingleObjectiveBenchmarkProblem:
-    return create_single_objective_problem_from_botorch(
-        test_problem_class=Branin,
-        test_problem_kwargs={},
-        lower_is_better=True,
-        num_trials=4,
-    )
-
-
 def get_single_objective_benchmark_problem(
     observe_noise_sd: bool = False,
     num_trials: int = 4,
     test_problem_kwargs: Optional[Dict[str, Any]] = None,
-) -> SingleObjectiveBenchmarkProblem:
+) -> BenchmarkProblem:
     return create_single_objective_problem_from_botorch(
         test_problem_class=Branin,
         test_problem_kwargs=test_problem_kwargs or {},