Revert D60472381 (facebook#2627)

Summary:
Pull Request resolved: facebook#2627

This diff reverts D60472381
Failure in OSS CI (T197664383), which will block other diffs in this repo. I plan to forward-fix by tomorrow or so, which is too slow.

Reviewed By: esantorella

Differential Revision: D60602573

fbshipit-source-id: 8da100be62a315582d14327bac396eba8c38b274

Author: generatedunixname89002005232357

ax/benchmark/tests/test_benchmark.py

@@ -6,11 +6,9 @@
 # pyre-strict
 
 import tempfile
-from typing import List
 from unittest.mock import patch
 
 import numpy as np
-import torch
 from ax.benchmark.benchmark import (
     _create_benchmark_experiment,
     benchmark_multiple_problems_methods,
@@ -23,10 +21,7 @@
     BenchmarkMethod,
     get_benchmark_scheduler_options,
 )
-from ax.benchmark.benchmark_problem import (
-    BenchmarkProblem,
-    create_single_objective_problem_from_botorch,
-)
+from ax.benchmark.benchmark_problem import create_single_objective_problem_from_botorch
 from ax.benchmark.benchmark_result import BenchmarkResult
 from ax.benchmark.methods.modular_botorch import get_sobol_botorch_modular_acquisition
 from ax.benchmark.metrics.base import GroundTruthMetricMixin
@@ -58,13 +53,6 @@
 from botorch.models.gp_regression import SingleTaskGP
 from botorch.optim.optimize import optimize_acqf
 from botorch.test_functions.synthetic import Branin
-from pyre_extensions import none_throws
-
-
-# Branin is subclassed here for convenience.
-class ThreeProblem(Branin):
-    def evaluate_true(self, X: torch.Tensor) -> torch.Tensor:
-        return torch.tensor(3.0, dtype=torch.float64)
 
 
 class TestBenchmark(TestCase):
@@ -102,84 +90,6 @@ def test_batch(self) -> None:
                 )
                 self.assertEqual(mock_optimize_acqf.call_args.kwargs["q"], batch_size)
 
-    def test_noisy_property_test(self) -> None:
-        problem = create_single_objective_problem_from_botorch(
-            test_problem_class=ThreeProblem,
-            test_problem_kwargs={"noise_std": 1.0},
-            lower_is_better=True,
-            num_trials=5,
-        )
-
-        res = benchmark_replication(
-            problem=problem,
-            method=get_sobol_benchmark_method(),
-            seed=0,
-        )
-        self.assertTrue((res.optimization_trace == 3).all())
-        df = none_throws(res.experiment).fetch_data().df
-        self.assertTrue(
-            (df.loc[df["metric_name"] == "ThreeProblem", "mean"] != 3).all()
-        )
-
-    def test_characterization_noisy_soo(self) -> None:
-        soo_problem = get_single_objective_benchmark_problem(
-            test_problem_kwargs={"noise_std": 1.0}
-        )
-        expected = [
-            104.36542659448965,
-            90.20076516036951,
-            56.18382601460175,
-            7.5867050504430775,
-        ]
-        self._noisy_characterization_test(problem=soo_problem, expected_result=expected)
-
-    def test_characterization_noisy_moo(self) -> None:
-        moo_problem = get_multi_objective_benchmark_problem(
-            test_problem_kwargs={"noise_std": 1.0}
-        )
-        expected = [0.0, 0.0, 6.952682254445629, 6.952682254445629]
-        self._noisy_characterization_test(problem=moo_problem, expected_result=expected)
-
-    def test_characterization_noisy_soo_surrogate(self) -> None:
-        # Surrogate construction shouldn't be random
-        torch.manual_seed(0)
-        soo_surrogate_problem = get_soo_surrogate(noise_stds=1.0)
-        expected = [
-            104.36542659448779,
-            104.36542659448779,
-            104.36542659448779,
-            104.36542659448779,
-        ]
-        self._noisy_characterization_test(
-            problem=soo_surrogate_problem,
-            expected_result=expected,
-        )
-
-    def test_characterization_noisy_moo_surrogate(self) -> None:
-        # Surrogate construction shouldn't be random
-        torch.manual_seed(0)
-        moo_surrogate_problem = get_moo_surrogate(noise_stds=1.0)
-        expected = [
-            24.999999999997655,
-            24.999999999997655,
-            24.999999999997655,
-            24.999999999997655,
-        ]
-        self._noisy_characterization_test(
-            problem=moo_surrogate_problem,
-            expected_result=expected,
-        )
-
-    def _noisy_characterization_test(
-        self, problem: BenchmarkProblem, expected_result: List[float]
-    ) -> None:
-        res = benchmark_replication(
-            problem=problem,
-            method=get_sobol_benchmark_method(),
-            seed=0,
-        )
-        self.assertTrue((res.optimization_trace == expected_result).all())
-
     def test_storage(self) -> None:
         problem = get_single_objective_benchmark_problem()
         res = benchmark_replication(

ax/utils/testing/benchmark_stubs.py

@@ -17,26 +17,24 @@
     MultiObjectiveBenchmarkProblem,
 )
 from ax.benchmark.benchmark_result import AggregatedBenchmarkResult, BenchmarkResult
-from ax.benchmark.metrics.benchmark import BenchmarkMetric
 from ax.benchmark.problems.surrogate import (
     MOOSurrogateBenchmarkProblem,
     SOOSurrogateBenchmarkProblem,
 )
 from ax.benchmark.runners.surrogate import SurrogateRunner
 from ax.core.experiment import Experiment
-from ax.core.objective import MultiObjective, Objective
 from ax.core.optimization_config import (
     MultiObjectiveOptimizationConfig,
     OptimizationConfig,
 )
-from ax.core.outcome_constraint import ObjectiveThreshold
 from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
 from ax.modelbridge.registry import Models
 from ax.modelbridge.torch import TorchModelBridge
 from ax.models.torch.botorch_modular.model import BoTorchModel
 from ax.models.torch.botorch_modular.surrogate import Surrogate
 from ax.service.scheduler import SchedulerOptions
 from ax.utils.common.constants import Keys
+from ax.utils.common.typeutils import checked_cast
 from ax.utils.testing.core_stubs import (
     get_branin_experiment,
     get_branin_experiment_with_multi_objective,
@@ -62,13 +60,11 @@ def get_single_objective_benchmark_problem(
 
 
 def get_multi_objective_benchmark_problem(
-    observe_noise_sd: bool = False,
-    num_trials: int = 4,
-    test_problem_kwargs: Optional[Dict[str, Any]] = None,
+    observe_noise_sd: bool = False, num_trials: int = 4
 ) -> MultiObjectiveBenchmarkProblem:
     return create_multi_objective_problem_from_botorch(
         test_problem_class=BraninCurrin,
-        test_problem_kwargs={} if test_problem_kwargs is None else test_problem_kwargs,
+        test_problem_kwargs={},
         num_trials=num_trials,
         observe_noise_sd=observe_noise_sd,
     )
@@ -98,97 +94,61 @@ def get_sobol_benchmark_method() -> BenchmarkMethod:
     )
 
 
-def get_soo_surrogate(noise_stds: float = 0.0) -> SOOSurrogateBenchmarkProblem:
-    outcome_name = "branin"
-    observe_noise_stds = True
+def get_soo_surrogate() -> SOOSurrogateBenchmarkProblem:
     experiment = get_branin_experiment(with_completed_trial=True)
-
-    optimization_config = OptimizationConfig(
-        objective=Objective(
-            metric=BenchmarkMetric(
-                name=outcome_name,
-                lower_is_better=False,
-                observe_noise_sd=observe_noise_stds,
-            ),
-            minimize=False,
-        )
-    )
-    surrogate = Surrogate(botorch_model_class=SingleTaskGP)
-    model_bridge = TorchModelBridge(
+    surrogate = TorchModelBridge(
         experiment=experiment,
         search_space=experiment.search_space,
-        model=BoTorchModel(surrogate=surrogate),
+        model=BoTorchModel(surrogate=Surrogate(botorch_model_class=SingleTaskGP)),
         data=experiment.lookup_data(),
         transforms=[],
     )
-    datasets = surrogate.training_data
     runner = SurrogateRunner(
-        name=outcome_name,
+        name="test",
         search_space=experiment.search_space,
-        outcome_names=[outcome_name],
-        get_surrogate_and_datasets=lambda: (model_bridge, datasets),
-        noise_stds=noise_stds,
+        outcome_names=["branin"],
+        get_surrogate_and_datasets=lambda: (surrogate, []),
     )
     return SOOSurrogateBenchmarkProblem(
         name="test",
         search_space=experiment.search_space,
-        optimization_config=optimization_config,
+        optimization_config=checked_cast(
+            OptimizationConfig, experiment.optimization_config
+        ),
         num_trials=6,
-        observe_noise_stds=observe_noise_stds,
+        observe_noise_stds=True,
         optimal_value=0.0,
         runner=runner,
         is_noiseless=runner.is_noiseless,
     )
 
 
-def get_moo_surrogate(noise_stds: float = 0.0) -> MOOSurrogateBenchmarkProblem:
-    observe_noise_stds = True
-    outcome_names = ["branin_a", "branin_b"]
-    # set this to be easy to beat, so hypervolume computations aren't all zero
-    ref_point = [10.0, 10.0]
-    surrogate = Surrogate(botorch_model_class=SingleTaskGP)
+def get_moo_surrogate() -> MOOSurrogateBenchmarkProblem:
     experiment = get_branin_experiment_with_multi_objective(with_completed_trial=True)
-    model_bridge = TorchModelBridge(
+    surrogate = TorchModelBridge(
         experiment=experiment,
         search_space=experiment.search_space,
-        model=BoTorchModel(surrogate=surrogate),
+        model=BoTorchModel(surrogate=Surrogate(botorch_model_class=SingleTaskGP)),
         data=experiment.lookup_data(),
         transforms=[],
     )
 
     runner = SurrogateRunner(
         name="test",
         search_space=experiment.search_space,
-        outcome_names=outcome_names,
-        get_surrogate_and_datasets=lambda: (model_bridge, surrogate.training_data),
-        noise_stds=noise_stds,
+        outcome_names=["branin_a", "branin_b"],
+        get_surrogate_and_datasets=lambda: (surrogate, []),
     )
-    metrics = [
-        BenchmarkMetric(
-            name=name,
-            lower_is_better=True,
-            observe_noise_sd=observe_noise_stds,
-        )
-        for name in outcome_names
-    ]
-    objectives = [Objective(metric=metric) for metric in metrics]
-    objective_thresholds = [
-        ObjectiveThreshold(metric=metric, bound=ref_p, relative=False)
-        for metric, ref_p in zip(metrics, ref_point)
-    ]
-    optimization_config = MultiObjectiveOptimizationConfig(
-        objective=MultiObjective(objectives=objectives),
-        objective_thresholds=objective_thresholds,
-    )
-
     return MOOSurrogateBenchmarkProblem(
         name="test",
         search_space=experiment.search_space,
-        optimization_config=optimization_config,
+        optimization_config=checked_cast(
+            MultiObjectiveOptimizationConfig, experiment.optimization_config
+        ),
         num_trials=10,
         observe_noise_stds=True,
         optimal_value=1.0,
-        reference_point=ref_point,
+        reference_point=[],
         runner=runner,
         is_noiseless=runner.is_noiseless,
     )